__all__ = [
    'Trajectory', 'datadir', 'D_arena', 'R_arena', 
    'draw_arena', 'arena_sample', 'rim_sample',
    'X_t', 'Y_t', 'speed_t', 'angle_t'
]

DEBUG = True

from __future__ import division, print_function

import os

import numpy as np
from numpy import load, array, arange, sqrt, median, pi
from numpy.random import rand
import matplotlib as mpl
import matplotlib.pyplot as plt

from tools.plot import quicktitle
from tools.radians import xy_to_rad_vec
from tools.filters import quick_boxcar
from tools.images import tiling_dims

from brian2.units import *
from brian2 import NeuronGroup, TimedArray

if DEBUG:
    %matplotlib inline

def logunit(v, name):
    if DEBUG:
        print('unit(%s) =' % name, get_unit(v))

datadir = '/Users/joe/source/phase_model/data'
posfile = os.path.join(datadir, 'pos.npy')
pos = load(posfile)

t_pos, x_pos, y_pos = pos
t_pos = pos[0] * second
traj_dt = median(t_pos[1:] - t_pos[:-1])

x_pos = pos[1] * cm
y_pos = pos[2] * cm

if DEBUG:
    print('dt =', traj_dt)
logunit(x_pos, 'x')
logunit(y_pos, 'y')

R_arena = 0.40 * meter
D_arena = 2 * R_arena

def draw_arena(ax=None):
    if ax is None:
        ax = plt.gca()
    ax.add_artist(mpl.patches.Circle((0,0), radius=R_arena / cm, fill=False, ec='0.3', lw=2, ls='dashed', zorder=-1))
    ax.axis(array([-1,1,-1,1]) * 1.2 * R_arena / cm)
    return ax

def randc(N):
    """Return a 2-row vector of uniform random samples on the unit disc
    """
    t = 2 * pi * rand(N)
    r = rand(N) + rand(N)
    r[r > 1] = 2 - r[r > 1]
    return r * cos(t), r * sin(t)

def arena_sample(n):
    unit_disc = randc(n)
    return R_arena * unit_disc[0], R_arena * unit_disc[1]

def test_arena_sample():
    xs, ys = arena_sample(100)
    plt.scatter(xs / cm, ys / cm, marker='.', color='b')
    plt.axis('equal')
    draw_arena();
    print('unit(xs) =', get_unit(xs))

def rim_sample(n):
    theta = np.linspace(0, 2*pi, n, endpoint=False)
    return R_arena * cos(theta), R_arena * sin(theta)

def test_rim_sample():
    xs, ys = rim_sample(10)
    plt.scatter(xs / cm, ys / cm, marker='.', color='b')
    plt.axis('equal')
    draw_arena();
    print('unit(xs) =', get_unit(xs))
    
if DEBUG:
    test_arena_sample()
    test_rim_sample()

# Position signal
X_t = TimedArray(x_pos, traj_dt)
Y_t = TimedArray(y_pos, traj_dt)

# Velocity signal
dx = quick_boxcar(x_pos[1:] - x_pos[:-1]) * meter
dy = quick_boxcar(y_pos[1:] - y_pos[:-1]) * meter

speed = (sqrt(dx**2 + dy**2) / traj_dt) / (cm / second) # speed doesn't need units
speed_t = TimedArray(speed, traj_dt)

angle = xy_to_rad_vec(dx, dy)
angle_t = TimedArray(angle, traj_dt)

logunit(speed, 'speed')
logunit(angle, 'angle')

def plot_traj_data():
    f, ax = plt.subplots(1, 3, figsize=(13,4))
    f.suptitle('trajectory data: (x,y), speed, angle')
    ax, sax, aax = ax
    
    ax.plot(X_t.values / cm, Y_t.values / cm, 'r-');
    draw_arena(ax)
    ax.axis('equal')
    ax.axis([-70,70,-50,50])
    ax.set(xlabel='x (cm)', ylabel='y (cm)')
    
    sax.plot(arange(speed.size) * traj_dt, speed)
    sax.set(xlabel='t (s)', ylabel='speed (cm/s)')
    
    aax.plot(arange(angle.size) * traj_dt, angle)
    aax.set(xlabel='t (s)', ylabel='angle (rad)')    
    
if DEBUG:
    plot_traj_data()

traj_eqs = """
    X = X_t(t) : meter
    Y = Y_t(t) : meter
    speed = speed_t(t) : 1  # cm /s
    angle = angle_t(t) : 1  # radians
"""

Trajectory = NeuronGroup(1, model=traj_eqs, dt=traj_dt, order=-1)

if DEBUG:
    print(traj_eqs)
    print(Trajectory)