import sys
sys.path.insert(0, 'W:/Lab/MERFISH_Data_060421/SingleCellRoutines')

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tqdm
from matplotlib.cm import get_cmap
import fbpca
import datetime
import seaborn as sns

from __init__scr import *
from __init__plots import *

# rotation transformation utility
def rotate(mat, vec_chiral):
    """(n,2) matrix as input
    """
    U, s, Vt = fbpca.pca(mat, k=2)
    coords = U.dot(np.diag(s))

    # det = -1 reflection; det = 1 rotation
    if np.linalg.det(Vt) < 0: # ~ -1
        coords[:,0] = -coords[:,0]
        Vt[0,:] = -Vt[0,:]
            
        delta = vec_chiral[coords[:,1]>0].sum() - vec_chiral[coords[:,1]<0].sum()
        if delta < 0: # rotate 180 
            coords = -coords
            Vt = -Vt
    
    return coords, Vt

# rotation main streamline
def rotation(DATA_DIR,samples):
    input = DATA_DIR+'/data/'+'processed_merfish_ad_mouse.hdf5'
    output_file = DATA_DIR+'/data/'+'processed_merfish_ad_mouse_rotated.hdf5'

    # read in data
    alldata = {}
    metas = {}
    datas = {}
    for sample in samples:
        # one sample
        meta = pd.read_hdf(input, 'meta_'+sample)
        meta['sample'] = sample 
        
        data = pd.read_hdf(input, 'mat_'+sample)
        
        toplot = meta.join(data)
        
        metas[sample] = meta
        datas[sample] = data
        alldata[sample] = toplot
    
    # perform rotation
    all_Vts = {}
    for sample in alldata.keys():
        toplot = alldata[sample]
        toplot[['x', 'y']], Vt = rotate(toplot[['center_x', 'center_y']].values, np.ones((toplot.shape[0],1)))
        
        all_Vts[sample] = Vt
        metas[sample][['x', 'y']] = toplot[['x', 'y']]
    
    # save alldata and rotation matrices as hdf5
    import h5py
    file = output_file
    for sample in alldata.keys():
        with h5py.File(file, 'a') as f:
            f['Vt_'+sample] = all_Vts[sample]
        metas[sample].to_hdf(file, 'meta_'+sample, 'a')
        datas[sample].to_hdf(file, 'mat_'+sample, 'a')
    
    with h5py.File(file, 'r') as f:
        print(f.keys())