# Some code to show distributions
import torch.distributions.studentT as studentT
import torch.distributions.laplace as laplace
import torch.distributions.normal as normal
import torch.distributions.uniform as uniform
import torch as torch
import matplotlib.pyplot as plt
import math

d1 = studentT.StudentT(torch.tensor([2.0]))
d2 = studentT.StudentT(torch.tensor([1.0]))
d3 = laplace.Laplace(torch.tensor([0.0]),torch.tensor([1/math.sqrt(2)]))
d4 = normal.Normal(torch.tensor([0.0]),torch.tensor([1.0]))
x = torch.linspace(-4.0, 4.0, 200)
# pdf
plt.plot(x,10**d1.log_prob(x))
plt.plot(x,10**d2.log_prob(x))
plt.plot(x,10**d3.log_prob(x))
plt.plot(x,10**d4.log_prob(x))
plt.title('different distributions (pdf)')
plt.legend(['studentT(df=2)','studentT(df=1)','laplace(0,1/sqrt(2))','normal(mean=0,var=1)'])
plt.show()

# log pdf
plt.plot(x,d1.log_prob(x))
plt.plot(x,d2.log_prob(x))
plt.plot(x,d3.log_prob(x))
plt.plot(x,d4.log_prob(x))
plt.title('different distributions (log pdf)')
plt.legend(['studentT(df=2)','studentT(df=1)','laplace(0,1/sqrt(2))','normal(mean=0,var=1)'])
plt.show()

# calculate covariance & correlation matrix using dataset
from sklearn.datasets import load_iris
iris = load_iris()
df_iris = torch.tensor(iris.data)
y = torch.tensor(iris.target)
columns=['sepal_length', 'sepal_width', 'petal_length', 'petal_width']

D = df_iris.shape[0]
diffs = df_iris-df_iris.mean(dim=0)
squares = (diffs)**2
variance_calc = torch.sum(squares,dim=0)/(D-1)
variance = torch.var(df_iris,dim=0)
cov_matrix_calc = torch.matmul(diffs.T,diffs)/(D-1)
print('----Covariance matrix(calulated & direct fn)-----')
print(cov_matrix_calc)
print(torch.cov(df_iris.T))
variance_from_cov_matrix = torch.unsqueeze(torch.diag(cov_matrix_calc),1)
corr_matrix_calc = cov_matrix_calc/torch.sqrt(variance_from_cov_matrix*variance_from_cov_matrix.T)
print('----Correlation matrix(calulated & direct fn)-----')
print(corr_matrix_calc)
print(torch.corrcoef(df_iris.T))

----Covariance matrix(calulated & direct fn)-----
tensor([[ 0.6857, -0.0424,  1.2743,  0.5163],
        [-0.0424,  0.1900, -0.3297, -0.1216],
        [ 1.2743, -0.3297,  3.1163,  1.2956],
        [ 0.5163, -0.1216,  1.2956,  0.5810]], dtype=torch.float64)
tensor([[ 0.6857, -0.0424,  1.2743,  0.5163],
        [-0.0424,  0.1900, -0.3297, -0.1216],
        [ 1.2743, -0.3297,  3.1163,  1.2956],
        [ 0.5163, -0.1216,  1.2956,  0.5810]], dtype=torch.float64)
----Correlation matrix(calulated & direct fn)-----
tensor([[ 1.0000, -0.1176,  0.8718,  0.8179],
        [-0.1176,  1.0000, -0.4284, -0.3661],
        [ 0.8718, -0.4284,  1.0000,  0.9629],
        [ 0.8179, -0.3661,  0.9629,  1.0000]], dtype=torch.float64)
tensor([[ 1.0000, -0.1176,  0.8718,  0.8179],
        [-0.1176,  1.0000, -0.4284, -0.3661],
        [ 0.8718, -0.4284,  1.0000,  0.9629],
        [ 0.8179, -0.3661,  0.9629,  1.0000]], dtype=torch.float64)