#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 23 16:49:55 2020

@author: malaikakironde
"""


#This code plots the average temperature of a network as a function of time,
#given that the heat diffuses using the heat equation

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import scipy as sp

def heat_eq (G,k,e0):
    L= nx.laplacian_matrix(G).todense() #create laplacian
    #print(L)
    temp_list=[]
    
    for t in np.arange(0,100,1):
        
        #time_list.append(t)
        A= -k*t*L
    
        E=sp.linalg.expm(A) #create exponential of matrix
        et= np.matmul(E,e0) #final state
    
        #Average temperature
        Av_temp = np.matrix.mean(et)
        temp_list.append(Av_temp)
    
    print(temp_list)
    t = np.arange(0,100,1)
    plt.plot(t,temp_list)
    plt.xlabel('Time')
    plt.ylabel('Average Temperature')
    plt.show()
    
    

def main():

    graph = input('Random or SWN? ').upper()
     
    if graph == 'RANDOM':
        n = int(input('Enter the number of nodes: '))
        m = int(input('Enter the number of edges: '))
        G = nx.gnm_random_graph(n, m) 
     
    elif graph == 'SWN':
        n = int(input('Enter the number of nodes in the graph: '))
        k = int(input('Enter the number of nearest neighbors that each node is connected to in ring topology: '))
        p = float(input('Enter the probabilty of rewiring each edge: '))
        G = nx.watts_strogatz_graph(n, k, p)  
             
    nx.draw(G, with_labels=True, node_color=range(n),cmap=plt.cm.Reds)
    plt.show()
     
    
    #Finding the inputs of the graph 
    k = eval(input('What is the heat constant? '))
     
    #finding the matrix for the initial temperature #order is equal to number of nodes
    e_vec = input("Enter the initial state as a vector, separated by';'(even laplacian) ")
    e0 = np.mat(e_vec)
     

    print(heat_eq(G,k,e0))

    
main()