import pandas as pd
import numpy as np
from sklearn.decomposition import PCA
import sys
import os
import matplotlib.pyplot as plt
import xarray as xr
import random
import math

# PCA cutoff procedure is heavily inspired by
# https://royalsocietypublishing.org/doi/10.1098/rsif.2014.0672 [Berman, 2014]

def moving_average(timedata, coldata, window_ms):
    # FIXME - for our datasets we need to consider timedata and use that for the convolution weights
    samplerate = 300                                # frames / second
    ms_per_frame = 1000.0 / samplerate              # ms / frame
    nframes = int(window_ms / ms_per_frame)         # ms / (ms/frame) = frames
    mavg = np.convolve(coldata, np.ones(nframes)/nframes, 'same')

    mavg[0:math.floor(nframes/2)] = mavg[math.floor(nframes/2)]
    mavg[-math.floor(nframes/2):] = mavg[-math.floor(nframes/2)-1]

    return mavg

def norm_df(df):
    return ((df - df.mean()) / df.std()).copy()

cattitles = {
    "angular_postural": "Posture (angles)",
    "distance": "Interpart distances",
    "velocity_egocentric": "Egocentric velocities",
    "egocentric_position": "Egocentric position",
    "velocity_relative": "Relative velocities",
}

def cat2title(cat):
    if cat in cattitles:
        return cattitles[cat]
    print(f"Unrecognized category {cat}")
    return cat

def getPlainDataCategories():
    return ["angular_postural", "distance", "velocity_egocentric", "egocentric_position"]

def getMovingAverageDataCategories():
    #return ["heading_velocity_cartesian", "velocity_egocentric", "velocity_relative"]
    return ["velocity_egocentric", "velocity_relative"]

def getColumnFeatureType(colname):
    if "behavior" in colname:
        return "behavior"
    if colname in ["vid_framenum", "vid_name", "frame", "time", "contig_framenum", "contig_num"]:
        return "sequence"

    if "absolute_position" in colname:
        if "_z" in colname:
            return "position_height"
        else:
            return "display"
    if "absolute_world_velocity_" in colname:
        return "display"

    if "vector_vector_angle" in colname:
        return "angular_postural"

    if "egocentric" in colname and "_pos_" in colname and "_xyangle_" in colname:
        return "angular_xy_part_pos"
    if "egocentric" in colname and "_relative_velocity_" in colname and "_xyangle_" in colname:
        return "angular_xy_part_relative_vel"

    if "vectorheading" in colname:
        if "velocity" in colname:
            return "heading_velocity_cartesian"
        else:
            return "heading_postural_cartesian"

    if "dist" in colname:
        return "distance"

    if "egocentric" in colname:
        if "velocity" in colname and "relative" in colname:
            return "velocity_relative"
        if "velocity" in colname:
            return "velocity_egocentric"
        if "_pos_" in colname:
            return "egocentric_position"

    if "closing_speed" in colname:
        return "velocity_relative"

    if "absolute_speed" in colname:
        return "velocity_world"

    print(f"**************** Unknown column {colname} *********************")
    return "other"

def isUsefulDataColumn(colname):
    category = getColumnFeatureType(colname)
    if category == "behavior" or category == "sequence":
        return False

    if category == "other":
        print(f"Warning: Unrecognized type for column {colname}")
        return False

    #"vector_vector_angle_an1_BodyCenter_HeadCenter_an1_SpineF_HeadCenter"]:
    if colname in ["egocentric_an1_pos_an1_SpineL_x", "egocentric_an1_pos_an1_SpineL_y",
                   "egocentric_an1_pos_an1_SpineL_z", "egocentric_an1_pos_an1_SpineM_y",
                   "egocentric_an1_vectorheading_an1_SpineL_x", "egocentric_an1_vectorheading_an1_SpineL_y",
                   "egocentric_an1_vectorheading_an1_SpineL_z", "vector_vector_angle_an1_BodyCenter_HeadCenter_an1_SpineF_HeadCenter"]: ## Figure out issue later!
        return False

    return True

def plotHist(pca, category, mode, outpath):
    fig, ax1 = plt.subplots(1, figsize=(10, 10))
    ax1.hist(pca.explained_variance_ratio_, bins=10)
    ax1.set_title(f'noise histo Plot - {category} - {mode}')
    plt.savefig(outpath)

def plotPCA(pca, category, mode, outpath):
    fig, (ax1, ax2) = plt.subplots(2, figsize=(10, 10))
    ax1.plot(np.arange(pca.n_components_) + 1,
             pca.explained_variance_ratio_, 'o-', linewidth=2, color='blue')
    ax1.set_title(f'Scree Plot - {cat2title(category)} - {mode}')
    ax1.set_xlabel('Principal Component')
    ax1.set_ylabel('Variance Explained')

    ax2.plot(np.arange(pca.n_components_) + 1,
             np.cumsum(pca.explained_variance_ratio_), 'o-', linewidth=2, color='blue')
    ax2.set_title(f'Cumulative variance explained by {cat2title(category)} PCs')
    ax2.set_xlabel('Principal Component')
    ax2.set_ylabel('Cumulative Variance Explained')
    plt.savefig(outpath)

def fitPCA(df, datacols):
    pca = PCA(n_components = min(40, len(datacols)))
    transformed = pca.fit_transform(df.loc[:,datacols])

    return transformed, pca

def calculatePCACutoff(pca_actual, pca_noise):
    # we'll call mean eigenvalue of the noisified PCA + 2 sigmas our cutoff for eigenvalues of our PCA
    threshold_eig = np.mean(pca_noise.explained_variance_) + np.std(pca_noise.explained_variance_) * 2
    last_accepted_pc = 0
    for i in range(pca_actual.n_components):
        if pca_actual.explained_variance_[i] > threshold_eig:
            last_accepted_pc = i
        else:
            print(f"rejected index {i}")
            break
    return threshold_eig, last_accepted_pc

def plotPCACutoff(pca_actual, pca_noise, threshold_eig, last_accepted_pc, category, outpath):
    fig, ax1 = plt.subplots(1, figsize=(10, 10))
    ax1.plot(np.arange(pca_actual.n_components_) + 1,
             pca_actual.explained_variance_ratio_, 'o-', linewidth=2, color='blue')
    ax1.plot(np.arange(pca_noise.n_components_) + 1,
             pca_noise.explained_variance_ratio_, 'o-', linewidth=2, color='red')
    #ax1.axhline(threshold_eig, linewidth=1, color='purple')
    ax1.axhline(pca_actual.explained_variance_ratio_[last_accepted_pc], linewidth=1, color='purple')
    ax1.axvline(last_accepted_pc+1, linewidth=1, color='purple')

    ax1.set_title(f"PC cutoff determination - {cat2title(category)}")
    ax1.set_xlabel("Component")
    ax1.set_ylabel("Variance Explained")
    ax1.legend(["Data", "Noise"])
    plt.savefig(outpath)

def importFiles(indir, category):
    filelist = []
    for f in os.listdir(indir):
        if os.path.isfile(os.path.join(indir, f)) and "features" in f and ".netcdf" in f:
            filelist.append(os.path.join(indir, f))

    df_accum = pd.DataFrame()
    for filename in filelist:
        print(f"Loading {filename}")
        ds = xr.open_dataset(filename)
        filtercols = list(filter(
                            lambda x: getColumnFeatureType(x) in ["sequence", "behavior", category],
                            ds.keys()))
        df = ds[filtercols].to_dataframe()
        df_accum = df_accum.append(df, ignore_index=True)
        ds.close()
        del df
        del ds

    datacols = list(filter(lambda x: getColumnFeatureType(x) == category, df_accum.columns))

    print("Norming...")
    df_normed_data = norm_df(df_accum.loc[:,datacols])

    print("Returning datasets...")
    df_accum.loc[:,datacols] = df_normed_data.loc[:,datacols]
    return df_accum

indir = "../features"
if len(sys.argv) >= 2:
    indir = sys.argv[1]

targetdir = os.path.join(indir, "PCs")
if len(sys.argv) >= 3:
    targetdir = sys.argv[2]

if not os.path.exists(targetdir):
    os.mkdir(targetdir)

for cat in getPlainDataCategories():
    print(f"Computing PCs for {cat}:")
    df = importFiles(indir, cat)
    pcaDf = pd.DataFrame()
    print(f"Constructed dataframe for {cat} [shape: {df.shape}]")
    print(df.shape)

    nondatacols = list(filter(lambda colname: not isUsefulDataColumn(colname), df.columns))
    datacols = list(filter(isUsefulDataColumn, df.columns))
    print(f"For category {cat}")
    print(f"Nondatacols: {nondatacols}")
    print(f"Datacols: {datacols}")

    for col in datacols:
        if np.any(np.isnan(df[col])):
            print(f"********* Col {col} has nans")

    # PCA of actual data
    transformed, pca_actual = fitPCA(df, datacols)
    plotPCA(pca_actual, cat, "data",
            os.path.join(targetdir, f"PCs-{cat}-data.png"))

    # PCA of noise [for determining how many PCs to keep]
    noisedf = df.copy()
    for col in noisedf.columns:
        noisedf[col] = np.random.permutation(noisedf[col])
    transformed_noise, pca_noise = fitPCA(noisedf, datacols)
    plotPCA(pca_noise, cat, "noise",
            os.path.join(targetdir, f"PCs-{cat}-noise.png"))

    plotHist(pca_noise, cat, "noisehist",
            os.path.join(targetdir, f"PCs-{cat}-noisehist.png"))

    threshold_eig, last_accepted_pc = calculatePCACutoff(pca_actual, pca_noise)
    plotPCACutoff(pca_actual, pca_noise, threshold_eig, last_accepted_pc, cat,
                  os.path.join(targetdir, f"PCs-{cat}-cutoff.png"))

    nToCopy = last_accepted_pc + 1
    for idx in range(nToCopy):
        pcaDf[f"PCA_{cat}_{idx}"] = transformed[:,idx]
    print(pcaDf)

    # add back the sequence and behavior columns
    for colname in list(filter(lambda x: getColumnFeatureType(x) in ["sequence", "behavior"], nondatacols)):
        pcaDf[colname] = df[colname]

    # convert float64 columns to float32 (for space)
    for colname in pcaDf.columns:
        if pcaDf[colname].dtype == np.float64:
            pcaDf[colname] = pcaDf[colname].astype(np.float32)

    outpath = os.path.join(targetdir, f"PCs-{cat}.netcdf")
    xr.Dataset.from_dataframe(pcaDf).to_netcdf(outpath, engine="h5netcdf")


    # Compute and save loading information

    print("Loadings:")
    loadings = pca_actual.components_.T * np.sqrt(pca_actual.explained_variance_)
    df_loadings = pd.DataFrame(loadings, columns=list(map(lambda x: f"PC_{x}", np.arange(len(pca_actual.components_)))), index=datacols)

    # Compute loadings only for the PCs we decided to keep based on the cutoff method above
    fh = open(os.path.join(targetdir, f"PCs-{cat}.loadings.csv"), "w")
    for pcnum in range(nToCopy):
        l = df_loadings[f"PC_{pcnum}"].values
        sorted_loading_indices = sorted(range(len(l)), key = lambda x: np.abs(l)[x], reverse=True)
        sorted_lvalues = l[sorted_loading_indices]

        line = f"PC{pcnum},"
        for ii in range(len(sorted_loading_indices)):
            line = line + f"{df_loadings.index[sorted_loading_indices[ii]]},{l[sorted_loading_indices[ii]]},"

        fh.write(line + "\n")

        #debugging
        #print(sorted_loading_indices[0:7])
        #print(sorted_lvalues[0:7])
        #print(line)
        #print(f"Top 7 for PC{pcnum}: ")
        #top7 = sorted_loading_indices[0:6]
        #print(df_loadings.index[top7])
    fh.close()

    del pca_actual
    del pca_noise

    del pcaDf
    del df

    plt.close('all')

##
#
#  Moving averages
#
##

for cat in getMovingAverageDataCategories():
    print(f"Computing moving averages for category {cat}")
    df = importFiles(indir, cat)
    print(f"Constructed dataframe for {cat} [shape: {df.shape}]")
    print(df.shape)

    nondatacols = list(filter(lambda colname: not isUsefulDataColumn(colname), df.columns))
    datacols = list(filter(isUsefulDataColumn, df.columns))
    print(datacols)
    print(f"Nondatacols: {nondatacols}")

    mavgDf = pd.DataFrame()
    for col in datacols:
        for period in [100, 150, 250, 500]:
            mavg = moving_average([], df[col].values, period)
            mavgDf[f"mavg_{period}_{col}"] = mavg
            # FIXME - this is a constant, that's not what we want.
            #mavgDf[f"sd_{period}_{col}"] = np.std(mavg)

    print("moving average cols:")
    print(mavgDf.columns)

    # PCA of actual data
    transformed, pca_actual = fitPCA(mavgDf, mavgDf.columns)
    plotPCA(pca_actual, cat, "data",
            os.path.join(targetdir, f"PCs-mavg-{cat}-data.png"))

    # PCA of noise [for determining how many PCs to keep]
    noisedf = mavgDf.copy()
    for col in noisedf.columns:
        noisedf[col] = np.random.permutation(noisedf[col])

    noisecols = mavgDf.columns
    transformed_noise, pca_noise = fitPCA(noisedf, noisecols)
    plotPCA(pca_noise, cat, "noise",
            os.path.join(targetdir, f"PCs-mavg-{cat}-noise.png"))

    plotHist(pca_noise, cat, "noisehist",
            os.path.join(targetdir, f"PCs-mavg-{cat}-noisehist.png"))

    threshold_eig, last_accepted_pc = calculatePCACutoff(pca_actual, pca_noise)
    plotPCACutoff(pca_actual, pca_noise, threshold_eig, last_accepted_pc, f"mavg_{cat}",
                  os.path.join(targetdir, f"PCs-mavg-{cat}-cutoff.png"))

    pcaDf = pd.DataFrame()
    nToCopy = last_accepted_pc + 1
    for idx in range(nToCopy):
        pcaDf[f"PCA_mavg_{cat}_{idx}"] = transformed[:,idx]
    print(pcaDf)

    # add back the sequence columns
    for colname in nondatacols:
        pcaDf[colname] = df[colname]

    # convert float64 columns to float32 (for space)
    for colname in pcaDf.columns:
        if pcaDf[colname].dtype == np.float64:
            pcaDf[colname] = pcaDf[colname].astype(np.float32)

    outpath = os.path.join(targetdir, f"PCs-mavg-{cat}.netcdf")
    xr.Dataset.from_dataframe(pcaDf).to_netcdf(outpath, engine="h5netcdf")

    plt.close('all')

sys.exit(1)