import numpy as np
import pandas as pd
import dendropy  as dp
import matplotlib as mpl

"""
Explanation: Introdução
Esse notebook traz as analises pedidas na disciplina Biologia Evolutiva - Bio507
Neste tutorial vamos utilizar a a linguagem Python 2.7 e os pacotes numpy, pandas, dendropy e matplotlib
Leitura de dados
End of explanation
"""


import numpy as np
import pandas as pd
import dendropy  as dp
import matplotlib as mpl

dados_brutos = pd.read_csv("./dados.csv")
num_traits = 4
traits = dados_brutos.columns[:num_traits]
dados_brutos.head(10)

"""
Explanation: Primeiro vamos ler os dados usando o pandas.
A função read_csv cria um data.frame com os dados.
End of explanation
"""


dados_brutos.describe()

"""
Explanation: Podemos usar a função describe para ter um resumo global, olhando para médias, desvio padrão e quartis.
End of explanation
"""


dados_brutos.groupby('ESPECIE').describe()

"""
Explanation: Podemos também usar a função groupby para aplicar funções em sub conjuntos dos dados, agrupando pela coluna ESPECIE
End of explanation
"""


medias = dados_brutos.groupby('ESPECIE').mean()
medias

std = dados_brutos.groupby('ESPECIE').std()
std

cv = std/medias
cv

"""
Explanation: Vamos calcular as médias, desvios padrão e coeficiantes de variação por caractere por espécie.
End of explanation
"""


# histogramas

# scatter plots

# boxplots

"""
Explanation: Visualizando os dados
Vamos usar o pacote matplotlib para fazer gráficos diagnósticos dos dados.
End of explanation
"""


cov_matrices = dados_brutos.groupby('ESPECIE').apply(lambda x: x.cov())
cor_matrices = dados_brutos.groupby('ESPECIE').apply(lambda x: x.corr())
especies_labels = list(pd.unique(dados_brutos['ESPECIE']))

cov_matrices

cor_matrices

"""
Explanation: Calculando matrizes de covariância
Vamos calcular as matrizes de covariância e correlação por espécie, além das médias, e agrupá-las num dicionário.
End of explanation
"""


cor_matrices.T['C']

"""
Explanation: Podemos também acessar as matrizes pelo nome da espécie:
End of explanation
"""


#heat plots

"""
Explanation: Ou fazer um heatplot de todas elas
End of explanation
"""


# Matrizes dentro e entre grupos

"""
Explanation: Analise de variância
Calculando as matrizes dentro de grupos, entre grupos e total
End of explanation
"""


# RandomSkewers

"""
Explanation: Comparando as matrizes
End of explanation
"""


tree = dp.Tree.get_from_string("(E, ((C, B)4,(A,D)3)2)1;", "newick")
tree.print_plot(display_width = 50, show_internal_node_labels = True, leaf_spacing_factor = 4)

"""
Explanation: Filogenia
Vamos incluir uma filogenia e calcular estados ancestrais
End of explanation
"""


get_node_name = lambda n: str(n.label or n.taxon or None)
nodes = [get_node_name(n) for n in tree.nodes()]

node_matrices = {}
node_sample_size = {}
for sp in especies_labels:
    new_matrix = np.array(cov_matrices.T[sp])
    node_matrices[sp] = new_matrix
    node_sample_size[sp] = dados_brutos[dados_brutos['ESPECIE'] == sp].shape[0]

# Tirando quem nao esta na filogenia e trocando os keys

node_means = {}

for sp in especies_labels:
    if tree.find_node_with_taxon_label(sp):
        new_key = get_node_name(tree.find_node_with_taxon_label(sp))
        node_means[new_key] = medias.T[sp]
        node_sample_size[new_key] = node_sample_size.pop(sp)
        node_matrices[new_key] = node_matrices.pop(sp)
    else:
        node_matrices.pop(sp)
        node_sample_size.pop(sp)

# Função que recebe uma lista de filhos e calcula a matriz, média e tamanho amostral pro ancestral

def ancestral_mean(child_labels):
    new_matrix = np.zeros((num_traits, num_traits))
    sample = 0
    new_mean = np.zeros(num_traits)
    for child in child_labels:
        node = get_node_name(child)
        new_matrix = new_matrix +\
            node_sample_size[node] * node_matrices[node]
        sample = sample + node_sample_size[node]
        new_mean = new_mean + node_sample_size[node] * node_means[node]
    new_matrix = new_matrix/sample
    new_mean = new_mean/sample
    return new_matrix, sample, new_mean

# Calculando as matrizes e tamanhos amostrais para todos os nós

for n in tree.postorder_node_iter():
    if get_node_name(n) not in node_matrices:
        node_matrices[get_node_name(n)], node_sample_size[get_node_name(n)], node_means[get_node_name(n)] = ancestral_mean(n.child_nodes())

"""
Explanation: Esse código usa as matrizes e médias calculadas anteriormente, junto com os tamanhos amostrais, para calcular valores ponderados para todos os nós da filogenia. 
A conta realizada para médias e matrizes é uma simples média ponderada.
End of explanation
"""


node_matrices

"""
Explanation: Isso resulta num dicionário, chamado node_matrices, com todas as matrizes para todos os nós.
End of explanation
"""


node_sample_size

"""
Explanation: Um dicionário de tamanhos amostrais
End of explanation
"""


node_means['1']

"""
Explanation: E um dicionário de médias.
End of explanation
"""


nodes

"""
Explanation: Temos tb uma lista de todos os nós:
End of explanation
"""


w_matrix

node_matrices['4']

"""
Explanation: É interessante notar como a matriz estimada para a raiz, ponderando as matrizes ao longo da filogenia, é idêntica à matriz ponderada dentro de grupos, calculada pelos modelos lineares anteriormente:
End of explanation
"""


delta_z = {}
beta = {}
for n in tree.nodes()[1:]: #começamos do 1 para pular a raiz, que não tem ancestral
    parent = get_node_name(n.parent_node)
    branch = get_node_name(n) + '_' + parent
    delta_z[branch] = node_means[get_node_name(n)] - node_means[parent]
    beta[branch] = np.linalg.solve(node_matrices[parent], delta_z[branch])

delta_z

beta

"""
Explanation: $\beta$ e $\Delta z$
Vamos agora estimar as mudanças evolutivas em cada ramo da filogenia, e, usando as matrizes ancestrais, calcular os gradientes de seleção estimados.
Os $\Delta z$ são apenas as diferenças nas médias de um nó com o seu acestral. Os $\beta$ são estimados com a equação de Lande:
$
\beta = G^{-1}\Delta z
$
End of explanation
"""


def vector_corr(x, y): return (np.dot(x, y)/(np.linalg.norm(x)*np.linalg.norm(y)))
corr_beta_delta_z = {}
for branch in delta_z:
    corr_beta_delta_z[branch] = vector_corr(beta[branch], delta_z[branch])

corr_beta_delta_z

"""
Explanation: Podemos agora calcular a correlação entre os $\beta$ e $\Delta z$
End of explanation
"""


corr_pc1 = {}
for branch in delta_z:
    parent = branch.split("_")[1]
    pc1 = np.linalg.eig(node_matrices[parent])[1][:,0]
    corr_pc1[branch] = abs(vector_corr(delta_z[branch], pc1))

corr_pc1

"""
Explanation: Podemos também calcular a relação entre a resposta evolutiva e o primeiro componente principal da matriz de covariação, a linha de menor resistência evolutiva. 
Como a direção primeiro componente é arbitrária, tomamos o valor absoluto da correlação.
End of explanation
"""


df_betas = pd.DataFrame.from_dict(beta, orient='index')
df_betas.columns = traits
df_betas

df_dz = pd.DataFrame.from_dict(delta_z, orient='index')
df_dz.columns = traits
df_dz

traits_c = list(traits)
traits_c.append('otu')
df_matrices = pd.DataFrame(columns=traits_c)
for node in node_matrices:
    df = pd.DataFrame(node_matrices[node], columns=traits, index = traits)
    df['otu'] = node
    df_matrices = df_matrices.append(df)
df_matrices


"""
Explanation: Podemos utilizar o pandas para formatar esses resultados em tabelas
End of explanation
"""


