from numpy.lib.function_base import angle
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import ticker
import glob
import os
from IPython import embed
import pycircstat as stats

from etrack import TrackingResult

fish_color = {'2019lepto03': {'main': 'red', 'dark': 'darkred', 'light': 'indianred'},
              '2020lepto20': {'main': 'orange','dark': 'sienna', 'light': 'sandybrown'},
              '2020lepto48': {'main': 'gold', 'dark': 'darkgoldenrod', 'light': 'yellow'},
              '2020lepto40': {'main': 'yellowgreen', 'dark': 'seagreen', 'light': 'mediumspringgreen'},
              '2020lepto42': {'main': 'turquoise', 'dark': 'darkslategrey', 'light': 'aquamarine'},
              '2020lepto43': {'main': 'cornflowerblue', 'dark': 'midnightblue', 'light': 'dodgerblue'},
              '2020lepto46': {'main': 'm', 'dark': 'darkmagenta', 'light': 'violet'}}

stimulus_angles = {1: 0, 2: 300 / 180 * np.pi, 3: 240/180 * np.pi, 4: np.pi, 5: 120/180 * np.pi, 6: 60/120 * np.pi}
light_conditions = ["light", "dark"]
stimulus_conditions = ["m", "f", "n"]

def plot_angles(angles, stimulus="", dark_light="", ax=None, color=None):
    if color is None:
        dark_color = "tab:blue"
    else:
        dark_color = color

    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1, polar=True)
    edges = np.linspace(0, 2*np.pi, 25)[:-1]
    n, bins, patches = ax.hist(angles / 180 * np.pi, bins=edges, density=True, align="left", color=dark_color, alpha=1.0 if dark_light == "dark" else 0.5)
    mid_angles = bins[1:] - np.diff(bins)/2
    mean_angle = stats.mean(angles/ 180 * np.pi)
    std_angle = stats.std(angles/ 180 * np.pi)

    ax.set_title(r"stim: %s, light: %s, $\sigma$=%.2f" %(stimulus, dark_light, std_angle/np.pi*180),
                 fontsize=7)
    ax.set_xticks(np.linspace(0.0, 2*np.pi, 7)[:-1])
    #ax.set_xticklabels(map(str, np.linspace(0, 300, 6)))
    ax.xaxis.set_major_formatter(ticker.FixedFormatter(list(map(lambda x:"%i$^\circ$" %x, np.linspace(0, 300, 6)))))
    ax.tick_params(which="major", labelsize=8)
    ax.set_yticklabels([])
    ymax = ax.get_ylim()[-1]
    ax.plot([mean_angle, mean_angle],[0.0, ymax], lw=1, ls="solid", color="red")
    if "n" not in stimulus:
        ax.scatter(0.0, ymax, marker="*", color="red", s=50, zorder=2)
    if ax is None:
        plt.show()


def get_trial_metadata(df, results_filename):
    conditions = get_trial_conditions(df, results_filename)
    stim_angle = None
    if not np.isnan(conditions.signal_output.values):
        stim_angle = stimulus_angles[int(conditions.signal_output)]
    release_angle = stimulus_angles[int(conditions.release)]
    subject = conditions.subject.values[0]
    stimulus_type = conditions.stimulus.values[0]
    light_condition = conditions.dark_light.values[0]
    return stim_angle, release_angle, subject, stimulus_type, light_condition


def get_trial_conditions(df, results_filename):
    movie_name = results_filename.split("/")[-1].split("DLC")[0] + ".mp4"
    trial_condition = df[df.dateiname == movie_name]
    return trial_condition


def get_angles(results_filename, min_likelihood=0.95):
    tr = TrackingResult(results_filename)
    angles = tr.angle_to_center("snout", min_likelihood=min_likelihood)
    return angles


def get_all_angles(df, subject, stimulus_condition="n", light_condition="dark", data_folder=os.path.join("..", "data"), min_likelihood=0.95, rotate_no_stim=False):
    subject_trials = df[df.subject == subject]
    trials = subject_trials[(subject_trials.stimulus == stimulus_condition) & (subject_trials.dark_light == light_condition)]
    all_angles = []
    all_angles_joined = []
    print(len(trials))
    for movie in trials.dateiname.values:
        results_files = glob.glob(os.path.join(data_folder, "%s%s"% (movie.split(".mp4")[0], "*.h5")))
        results_file = results_files[0] if len(results_files) > 0 else None
        if results_file is None:
            print("Could not find tracking results for movie %s!" % movie)
            continue
        angles = get_angles(results_file, min_likelihood=min_likelihood)
        if len(angles) == 0:
            print("Skipping movie %s, no data!" % movie)
            continue
        stimulus_angle, release_angle, _, _, _ = get_trial_metadata(df, results_file)
        if stimulus_condition != "n":
            if stimulus_angle is None:
                print("Stimulus angle is None for movie %s, stimulus %s" % (movie, stimulus_condition))
                continue
            angles += (360 - stimulus_angle)
            angles[angles>360] -= 360
        else:
            if rotate_no_stim:
                angles += (360 - release_angle)
                angles[angles>360] -= 360

        all_angles.append(angles)
        all_angles_joined.extend(angles)
        
    return all_angles, all_angles_joined


def characterize_behavior(phis):
    """[summary]

    Zugewandt =  3/2 π   <   𝞅   ≤   π/2 
    Zugewandt =  270   <   𝞅   ≤   90 
    
    Abgewandt =  π/2  <  𝞅   ≤   3/2 π
    Abgewandt =  90  <  𝞅   ≤   270

    Interagierend = 7/8 π  <  𝞅  ≤  1/4 π
    Interagierend = 315  <  𝞅   ≤ 45

    Vermeiden = 3/4 π <  𝞅   ≤  5/4 π
    Vermeiden = 135 <  𝞅   ≤  225
    """
    attracted = []
    repelled = []
    interacting = []
    avoiding = []

    for i in range(len(phis)):
        phi = phis[i]
        attracted.append(len(phi[(phi > 270) | (phi <= 90)]) / len(phi))
        repelled.append(len(phi[(phi > 90) & (phi <= 270)]) / len(phi))
        interacting.append(len(phi[(phi > 315) | (phi <= 45)]) / len(phi))
        avoiding.append(len(phi[(phi > 135) & (phi <= 225)]) / len(phi))

    return attracted, repelled, interacting, avoiding


def characterize_behavior_all(phis):
    attracted = len(phis[(phis > 270) | (phis <= 90)]) / len(phis)
    repelled = len(phis[(phis > 90) & (phis <= 270)]) / len(phis)
    interacting = len(phis[(phis > 315) | (phis <= 45)]) / len(phis)
    avoiding = len(phis[(phis > 135) & (phis <= 225)]) / len(phis)

    return attracted, repelled, interacting, avoiding

if __name__ == "__main__":
    df = pd.read_csv("../data/Tuebingen_tabelle_fertig.csv", sep=",", index_col=0)
    files = glob.glob("../data/*.h5")
    
    for f in files[:1]:
        angles = get_angles(f)
        stimulus_angle, release_angle, subject, stimulus_condition, light_condition = get_trial_metadata(df, f)
    
    subjects = df.subject.unique()
    subplot_labels = ["A", "B", "C"]
    results = []
    for subject in subjects:
        fig = plt.figure(figsize=(7, 5))
        for i, light in enumerate(light_conditions):
            for j, stim in enumerate(stimulus_conditions):
                print(subject, light, stim)
                axis = fig.add_subplot(len(light_conditions), len(stimulus_conditions), i * len(stimulus_conditions) + j + 1, polar=True)
                all, all_joined = get_all_angles(df, subject, stimulus_condition=stim, light_condition=light, min_likelihood=0.85, rotate_no_stim=True)

                if len(all_joined) > 0:
                    attracted, repelled, interacting, avoiding = characterize_behavior_all(np.array(all_joined))
                    #attracted, repelled, interacting, avoiding = characterize_behavior(all)

                    results.append({"subject": subject, "light_condition": light, "stimulus": stim, "attracted": attracted, 
                                    "repelled": repelled, "interacting": interacting, "avoiding": avoiding})
                    plot_angles(np.array(all_joined), stim, light, ax=axis, color=fish_color[subject]["dark"])
                if i == 0:
                    axis.text(-0.2, 1.15, subplot_labels[j], transform=axis.transAxes, fontsize=12, fontweight="bold")
        fig.subplots_adjust(wspace=0.45, hspace=0.25, left=0.07, top=0.925, right=0.975, bottom=0.05)
        fig.savefig("%s_preferred_angles.pdf" % subject)
        plt.close()
    results_df = pd.DataFrame(results)
    results_df.to_csv("behavior_results_simple.csv", sep=";")
    embed()

# does the fish go to the previous stimulus location?