# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Data Processing

# # Setup

import random,copy,math,time,os,csv,sys
import scipy.io as sio # for loading .mat files
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

DATA_DIRECTORY = './../../experimental_data/raw_data/All-datasets - Oct2020/'

data_dict = sio.loadmat('./../../experimental_data/raw_data/All-datasets - Oct2020/CGG-SpCas9.mat')
data_dict.keys()

# ## Load raw data

data_raw = []
for filename in os.listdir(DATA_DIRECTORY):
        if filename.endswith(".mat"):
            experiment = filename[:-4]
            print(experiment)
            data_dict = sio.loadmat(DATA_DIRECTORY +filename)
            if 't' in data_dict:
                data_dict["experiment"] = experiment
                if experiment == 'NoGuidedRNA':
                    data_dict["experiment"] = 'Control'
                if '-' in experiment:
                    defect = experiment.split('-')[0]
                    if defect in 'NonRepeatedSequence':
                        data_dict['defect'] = 'NRS'
                    else:
                        data_dict['defect'] = experiment.split('-')[0]
                    data_dict['nuclease'] = experiment.split('-')[1]

    
                if 'AllCells' in data_dict:
                    data_dict['AllCellsBF'] = data_dict['AllCells']
                    #data_dict['t'] = np.array(range(len(data_dict['AllCells'])))
                #else:
                    #data_dict['t'] = np.array(range(len(data_dict['AllCellsBF'])))
                if 'AllGFPCells' in data_dict:
                    data_dict['AllCellsGFP'] = data_dict['AllGFPCells']
                data_raw.append(data_dict)
                print('  success!')
            else:
                print('  missing data!')


def make_single_cell_dataframe(dct):
    '''
    convert the dictionary of raw data to a data frame
    '''
    data = pd.DataFrame() # dataframe for single-cell data
    
    
    n_wells = len(dct['AllCellsBF'][0])
    n_points = len(dct['t'][0])
    data['well'] = np.concatenate([np.ones(n_points)*j for j in range(n_wells)])
    data['bf'] = dct['AllCellsBF'].T.reshape(n_wells*n_points)
    
    if 'AllCellsGFP' in dct:
        data['gfp'] = dct['AllCellsGFP'].T.reshape(n_wells*n_points)
    else: # this means there is no defect
        data['gfp'] = np.nan
        
    data['time'] = np.concatenate([dct['t'][0] for j in range(n_wells)])
    data['experiment'] = dct['experiment']
    
    if 'defect' in dct:
        data['defect'] = dct['defect']
    else:
        data['defect'] = 'none'
    if 'nuclease' in dct:
        data['nuclease'] = dct['nuclease']
    else:
        data['nuclease'] = 'none'
    return data

# +
data = pd.concat([make_single_cell_dataframe(dct) for dct in data_raw])


data['initial_cells'] = initial_cells
data.to_csv('./../../experimental_data/processed_data/single_cell_data.csv')

# +

data_avg =data.groupby(['experiment','time','nuclease','defect']).mean()
data_avg.to_csv('./../../experimental_data/processed_data/avg_data.csv')


data_avg =data[data.initial_cells==1].groupby(['experiment','time','nuclease','defect']).mean()
data_avg.to_csv('./../../experimental_data/processed_data/avg_data_one_cell.csv')
# -

[plt.plot(data.groupby(['experiment','time']).mean().bf[exp].values,'C0-') \
 for exp in data.experiment.unique()];
[plt.semilogy(data.groupby(['experiment','time']).mean().gfp[exp].values,'C1--') \
for exp in data.experiment.unique()];

data.groupby(['experiment','time']).mean()

initial_cells =[]
for exp in data.experiment.unique():
    df = data[data.experiment == exp]
    initial_cells.append(np.array([df[df.well==well].bf.values[0] for well in df.well.values]))
initial_cells= np.concatenate(initial_cells)

plt.plot([len(data[data.initial_cells==k])/len(data) for k in [1,2,3,4,5,6]],'o-')

data



StarcoderdataJupyterScriptsDedupFiltered