# -*- coding: utf-8 -*-
"""
Created on Tue Sep 22 12:02:07 2020

@author: annak
"""

import matplotlib.pyplot as plt 
import numpy as np

# Description of System (Hamiltonian, initial population, master equation)

H = np.array([[-2, 5, 10], [2, -12, 0], [0, 7, -10]])

p0 = np.array([1,1,1])
q0 = np.array([2,0,1])

master = lambda t0, v: np.dot(H, v) #ODE governing the closed system

T = 0.5
N = 30

n = 3 #Number of states

# Solution of ODE using Expliziter Euler :) (Possibly to be investigated with other ODE solvers)

def integrate(method, rhs, y0, T, N):
    y = np.empty((N+1,) + y0.shape)

    t0, dt = 0.0, T/N
    y[0,...] = y0
    for i in range(0, N):
        y[i+1,...] = method(rhs, y[i,...], t0 + i*dt, dt)

    t = np.arange(N+1)*dt
    return t, y

def explicit_euler_step(rhs, y0, t0, dt):
    return y0 + dt*rhs(t0, y0)

def explicit_euler(rhs, y0, T, N):
    return integrate(explicit_euler_step, rhs, y0, T, N)

# Calculation of Entropy

def rel_ent(p, q):
    S = 0
    for i in range(1, n):
        S += p[i]*np.log(p[i]/q[i])
    return S

# Plotting of Entropy

t, p = explicit_euler(master, p0, T, N)
t, q = explicit_euler(master, q0, T, N)

#plt.figure(figsize=(20,10))
#plt.plot(t, p)
#plt.plot(t, q)

S = np.zeros(N)
for i in range(1, N):
    S[i] = rel_ent(p[i,:], q[i,:])
   
plt.figure(figsize=(20,10))
plt.grid()
plt.plot(t[1:N], S[1:N])
plt.show()