
from scipy.spatial import KDTree
from scipy.spatial import Voronoi
import numpy as np
import pandas as pd
from nufeb_tools import __version__
# TODO modify to account for edge effects
def fitness_metrics(obj):
    """
    Function to calculate colony-level fitness metrics. 

    Mother cell: Cell ID which seeded the colony

    Type: Cell type - cyanobacteria are type 1  and E. coli are type 2

    Voronoi area: Area of species-specific Voronoi Tesselation at the beginning of the simulation

    IPTG: Sucrose induction level

    total biomass: Biomass of each colony at the end of the simulation (fg)

    Nearest 1: Distance to nearest cyanobacteria colony

    Nearest 2: Distance to nearest E. coli colony 

    Nearest Neighbor: Distance to nearest colony

    IC1: Average distance to nearest cyanobacteria colony

    IC2: Average distance to nearest E. coli colony

    IC: Average distance to nearest colony

    Relative Neighbor Dist 1: Distance to nearest cyanobacteria colony divided by IC1

    Relative Neighbor Dist 2: Distance to nearest E. coli colony divided by IC2

    Relative Neighbor Dist: Distance to nearest colony divided by IC

    Z1: Relative neighbor distance 1 divided by sqrt(D_sucrose/mu_cyano)

    Z2: Relative neighbor distance 2 divided by sqrt(D_sucrose/mu_ecw)

    Z1_2: Relative neighbor distance 1 divided by sqrt(D_sucrose/mu_ecw)

    Z2_1: Relative neighbor distance 2 divided by sqrt(D_sucrose/mu_cyano)

    LogNearest 1: log(Nearest 1)

    LogNearest 2: log(Nearest 2)

    LogNearest: log(Nearest Neighbor)

    Inv1: Inverse sum of neighbor distance 1

    Inv2: Inverse sum of neighbor distance 2

    Log Inv1: Log squared inverse sum of neighbor distance 1

    Log Inv2: Log squared inverse sum of neighbor distance 2
    
    Colony Area: 2D area of colony at the end of the simulation

    Args:
        obj (nufeb_tools.utils.get_data): 
            Data object collected with nufeb_tools.utils.get_data

    Returns:
        pandas.DataFrame: 
            Dataframe containing colony number (mother cell ID), cell type, total biomass, colony area, Voronoi area, nearest neighbor, mean neighbor distance, etc.
    """
    if not hasattr(obj,'colonies'):
        obj.get_mothers()
    obj.count_colony_area(obj.Timesteps[-1])
    D_suc = obj.metadata['Diff_c']['suc']
    mu_ecw = obj.metadata['ecw']['GrowthRate']
    mu_cy = obj.metadata['cyano']['GrowthRate']
    df = obj.colonies.copy()
    # calculate voronoi area
    dfs = list()
    for type_ in df.type.unique():
        IDs = df[(df.Timestep ==0) & (df.type == type_)][['mother_cell','type']]
        points = df[(df.Timestep ==0) & (df.type == type_)][['x','y']].values
        vor = Voronoi(points)
        areas = [abs(np.sum( [0.5, -0.5] * vor.vertices[vor.regions[i]] * np.roll( np.roll(vor.vertices[vor.regions[i]], 1, axis=0), 1, axis=1) )) for i in range(len(vor.regions))]
        IDs.loc[:,'Voronoi Area'] = areas[1:]
        dfs.append(IDs)
    metrics = pd.concat(dfs)
    metrics.loc[:,'IPTG'] = obj.IPTG
    metrics.loc[:,'Time'] = df.time.iloc[-1]
    x_max = obj.metadata['Dimensions'][0]
    y_max = obj.metadata['Dimensions'][1]
    metrics.loc[:,'Distance from center']=df.loc[(df.Timestep==0)].apply(lambda x:np.sqrt((x_max/2-x['x'])**2 +(y_max/2-x['y'])**2),axis=1)
    # total biomass per colony
    biomasses = df[df.Timestep==obj.Timesteps[-1]].groupby('mother_cell').sum().reset_index()[['mother_cell','biomass']]
    biomasses.columns=['mother_cell','total biomass']
    initial_biomass = df[df.Timestep==0].groupby('mother_cell').sum().reset_index()[['mother_cell','biomass']]
    initial_biomass.columns=['mother_cell','initial biomass']
    # Calculate nearest neighbors
    df3 = df[df.Timestep == 0]
    arr = df3[['x','y','z']].to_numpy()
    tree = KDTree(arr)
    d, i = tree.query(arr, k=2)
    arr1 = df3[df3.type==1][['x','y','z']].to_numpy()
    tree1 = KDTree(arr1)
    d1, i1 = tree1.query(df3[['x','y','z']].to_numpy(), k=2)
    arr2 = df3[df3.type==2][['x','y','z']].to_numpy()
    tree2 = KDTree(arr2)
    d2, i2 = tree2.query(df3[['x','y','z']].to_numpy(), k=2)
    n1 = list()
    n2 = list()
    for i in range(len(d1)):
        if d1[i,0]==0:
            n1.append(d1[i,1])
        else:
            n1.append(d1[i,0])
    for i in range(len(d2)):
        if d2[i,0]==0:
            n2.append(d2[i,1])
        else:
            n2.append(d2[i,0])
    df3.loc[:,'Nearest 1']=n1
    df3.loc[:,'Nearest 2']=n2
    df3.loc[:,'Nearest Neighbor'] = d[:,1]
    df3 = df3[['mother_cell','Nearest 1','Nearest 2','Nearest Neighbor']]
    df3.loc[:,'IC1'] = df3.loc[:,'Nearest 1'].mean()
    df3.loc[:,'IC2'] = df3.loc[:,'Nearest 2'].mean()
    df3.loc[:,'IC'] = df3.loc[:,'Nearest Neighbor'].mean()
    df3.loc[:,'Relative Neighbor Dist 1'] = df3.loc[:,'Nearest 1']/df3.loc[:,'IC1']
    df3.loc[:,'Relative Neighbor Dist 2'] = df3.loc[:,'Nearest 2']/df3.loc[:,'IC2']
    df3.loc[:,'Relative Neighbor Dist'] = df3.loc[:,'Nearest Neighbor']/df3.loc[:,'IC']
    df3.loc[:,'Z1'] = df3.loc[:,'Relative Neighbor Dist 1']/np.sqrt(D_suc/mu_cy)
    df3.loc[:,'Z2'] = df3.loc[:,'Relative Neighbor Dist 2']/np.sqrt(D_suc/mu_ecw)
    df3.loc[:,'Z1_2'] = df3.loc[:,'Relative Neighbor Dist 1']/np.sqrt(D_suc/mu_ecw)
    df3.loc[:,'Z2_1'] = df3.loc[:,'Relative Neighbor Dist 2']/np.sqrt(D_suc/mu_cy)
    df3.loc[:,'LogNearest 1'] = np.log(df3['Nearest 1'])
    df3.loc[:,'LogNearest 2'] = np.log(df3['Nearest 2'])
    df3.loc[:,'LogNearest'] = np.log(df3.loc[:,'Nearest Neighbor'])
    inv1 = list()
    for i in df[df.Timestep == 0][['x','y','z']].to_numpy():
        d1, i1 = tree1.query(i,k=2)
        if d1[0]==0:
            inv1.append(np.sum(1/d1[1]))
        else:
            inv1.append(np.sum(1/d1[0]))
    inv2 = list()
    inv3 = list()
    for i in df[df.Timestep == 0][['x','y','z']].to_numpy():
        d2, i2 = tree2.query(i,k=2)
        if d2[0]==0:
            inv2.append(np.sum(1/d2[1]))
        else:
            inv2.append(np.sum(1/d2[0]))
    # Calculate log inverse squared neighbor distance
    log_inv1 = list()
    for i in df[df.Timestep == 0][['x','y','z']].to_numpy():
        d1, i1 = tree1.query(i,k=2)
        if d1[0]==0:
            log_inv1.append(np.log(np.sum(1/(d1[1]**2))))
        else:
            log_inv1.append(np.log(np.sum(1/(d1[0]**2))))
    log_inv2 = list()
    for i in df[df.Timestep == 0][['x','y','z']].to_numpy():
        d2, i2 = tree2.query(i,k=2)
        if d2[0]==0:
            log_inv2.append(np.log(np.sum(1/(d2[1]**2))))
        else:
            log_inv2.append(np.log(np.sum(1/(d2[0]**2))))

    df3.loc[:,'Inv1']=inv1
    df3.loc[:,'Inv2']=inv2
    df3.loc[:,'Log Inv1']=log_inv1
    df3.loc[:,'Log Inv2']=log_inv2

    colony_area = df[df.Timestep==obj.Timesteps[-1]][['mother_cell','Colony Area']].drop_duplicates()

    metrics = metrics.merge(biomasses,on='mother_cell')
    metrics = metrics.merge(df3,on='mother_cell')
    metrics = metrics.merge(colony_area,on='mother_cell')
    metrics = metrics.merge(initial_biomass,on='mother_cell')
    return metrics
