#!/usr/bin/env python
"""
Created on Thu Jul 16 15:08:27 2020

@author: Aloma Blanch
"""


import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from statistics import mean 
# from scipy.signal import find_peaks
import matplotlib.transforms as mtransforms


def error_plot(folder,t_step,r_criteria,save_path):
    # Load iterations and residual error
    histor = folder + '/histor.dat'
    input = open(histor, 'r')
    output = open(folder + "/histor_better.dat", 'w')
    output.writelines(line.strip() +'\n' for line in input)
    input.close()
    output.close()
    error_info = pd.read_csv(folder + "/histor_better.dat", sep=' ', header=None, usecols=(0,1,2))

    # Select only last iteration of residual error
    error=[] 
    for n in range(1,error_info.shape[0]):
        if (error_info[0][n]>error_info[0][n-1]):
            error.append(error_info[2][n-1])

    time = np.linspace(start = t_step, stop = len(error)*t_step, num = len(error))

    # Liniar Scale
    fig, ax = plt.subplots()
    ax.plot(time,error)
    ax.plot(time,r_criteria*np.ones(len(error)),'r')
    ax.set(xlabel='Time steps', ylabel='Residual error',
        title='Last nonlinear residual error for each time step')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    # plt.show()
    plt.savefig(save_path + '/Last_nonlin_res_error.pdf')
    # Semilog scale
    fig, ax = plt.subplots()
    ax.semilogy(time,error)
    ax.semilogy(time,r_criteria*np.ones(len(error)),'r')
    ax.set(xlabel='Time steps', ylabel='Residual error',
        title='Log - Last nonlinear residual error for each time step')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    # plt.show()
    plt.savefig(save_path + '/Log_Last_nonlin_res_error.pdf')
    fig.savefig(save_path + '/Log_Last_nonlin_res_error.jpg',dpi=150)


def periodicity(project,folder,dt,T_cyc,n_cyc,save_path):
    pressure = np.loadtxt(folder+'/PHistRCR.dat',skiprows=2,)
    time = np.linspace(0,T_cyc*n_cyc,round(T_cyc/dt*n_cyc))
    peak_P = []
    peak_P_pos = []
    Nc = round(T_cyc/dt)
    for i in range(0,n_cyc):
        peak_P.append(np.amax(pressure[i*Nc:Nc*(i+1),-1])/1333.22)
        peak_P_pos.append(np.where(pressure[i*Nc:Nc*(i+1),-1] == np.amax(pressure[i*Nc:Nc*(i+1),-1]))[0][0]+Nc*i)
    
    peak_Pdiff = [peak_P[n]-peak_P[n-1] for n in range(1,len(peak_P))]
    peak_Pdiff = list(map(abs, peak_Pdiff)) 
    
    fig, ax = plt.subplots()
    ax.plot(time,pressure[:,-1]/1333.22,'b')
    ax.plot(time[peak_P_pos], peak_P,'ro',label='Cylce pike')
    ax.set(xlabel='time [s]', ylabel='Pressure [mmHg]',
        title='Pressure at last outlet')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.legend(loc=0)
    # plt.show()
    plt.savefig(save_path + '/periodicity.pdf')
    fig.savefig(save_path + '/periodicity.jpg',dpi=150)
    
    if (peak_Pdiff[-1]<=1): 
        print('The numerical simulation \'{0}\' has achieve periodicity!\nSystolic Blood Pressure (SBP):\nsecond-last cycle = {1:.2f} mmHg,\nlast cycle = {2:.2f} mmHg,\n\u0394mmHg = {3:.2f} mmHg'.format(project,peak_P[-2],peak_P[-1],peak_Pdiff[-1]))  
        # txt = ['The numerical simulation \'{0}\' has achieve periodicity!'.format(project), 'Systolic Blood Pressure (SBP):', 'Second-last cycle = {0:.2f} mmHg'.format(peak_P[-2]), 'Last cycle = {0:.2f} mmHg'.format(peak_P[-1]), 'Delta_mmHg = {0:.2f} mmHg'.format(peak_Pdiff[-1])]
        txt = ['The numerical simulation \'{0}\' has achieve periodicity!'.format(project), 'Systolic Blood Pressure (SBP):','Second-last cycle = {0:.2f} mmHg - Last cycle = {1:.2f} mmHg - Delta_mmHg = {2:.2f} mmHg'.format(peak_P[-2],peak_P[-1],peak_Pdiff[-1])]
    return txt
    
def pressure(folder,N_ts,T_cyc,dt,n_cyc,save_path):
    pressure = np.loadtxt(folder+'/PHistRCR.dat',skiprows=2,)
    Nc = round(T_cyc/dt)
    time = np.linspace(0,T_cyc,Nc)
    fig, ax = plt.subplots()
    SBP = np.empty(pressure.shape[1])
    DBP = np.empty(pressure.shape[1])
    MBP = np.empty(pressure.shape[1])
    for i in range(0,pressure.shape[1]):
        ax.plot(time,pressure[N_ts-Nc:N_ts,i]/1333.22,label='ROI-'+str(i+2)) 
        SBP[i] = (np.amax(pressure[N_ts-Nc:N_ts,i]/1333.22))
        DBP[i] = (np.amin(pressure[N_ts-Nc:N_ts,i]/1333.22))
        MBP[i] = (mean(pressure[N_ts-Nc:N_ts,i]/1333.22)) 
    PP = SBP-DBP    
    ax.set(xlabel='time [s]', ylabel='Pressure [mmHg]',
        title='Pressure at each outlet')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.legend(loc=0)
    # plt.show()
    plt.savefig(save_path + '/pressure.pdf')
    fig.savefig(save_path + '/pressure.jpg',dpi=150)
    return (DBP,MBP,SBP,PP)


def flow(folder,N_ts,T_cyc,dt,n_cyc,save_path):
    flow = np.loadtxt(folder+'/QHistRCR.dat',skiprows=2,)
    Nc = round(T_cyc/dt)
    time = np.linspace(0,T_cyc,Nc)
    fig, ax = plt.subplots()
    Q = np.empty(flow.shape[1])
    for i in range(0,flow.shape[1]):
        ax.plot(time,flow[N_ts-Nc:N_ts,i],label='ROI-'+str(i+2))
        Q[i] = (mean(flow[N_ts-Nc:N_ts,i]))
   
    ax.set(xlabel='time [s]', ylabel='Flow [mL/s]',
        title='Flow at each outlet')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.legend(loc=0)
    # plt.show()
    plt.savefig(save_path + '/flow.pdf')
    fig.savefig(save_path + '/flow.jpg',dpi=150)
    return Q
    
def inlet_flow_waveform(project_folder,t_btw_rst,N_ts,dt,T_cyc,n_cyc,save_path):
    x = np.loadtxt(project_folder+'/ROI-1.flow')
    t = x[:,0]
    Q = -x[:,1]
    Nt_pts = np.linspace(t_btw_rst,N_ts,int(N_ts/t_btw_rst))
    t_pts = Nt_pts*dt
    # Put all the time values on a single cardiac cylce            
    for n in range(len(t_pts)):
        tmp=divmod(t_pts[n],T_cyc)
        t_pts[n]=tmp[1]
        if round(tmp[1],3) ==  0:
            t_pts[n]=T_cyc
    # Interpolate the flow rate to obtain the location of the point
    Q_pts = np.interp(t_pts, t, Q)

    fig, ax = plt.subplots()
    ax.plot(t, Q, 'r')
    ax.plot(t_pts, Q_pts, 'ob',label='Time steps saved')
    trans_offset = mtransforms.offset_copy(ax.transData, fig=fig,
                                       x=-0, y=0.15, units='inches')
    ax.set(xlabel='Time [s]', ylabel='Flow Rate - Q [mL/s]',
       title='Inlet Flow Rate Waveform')
    ax.set_ylim([-10, 90])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    # Adding label to the points
    
    time = []
    for i in range(0,np.unique(np.round(t_pts,3)).shape[0]):
        time.append('$t_'+str(i+1)+'$')
    for x, y, t in zip(t_pts[(-n_cyc-1):], Q_pts[(-n_cyc-1):], time):
        plt.text(x, y, t, transform=trans_offset, fontsize=12)    
    ax.legend(loc=0)
    # plt.show()
    plt.savefig(save_path + '/inlet_waveform.pdf')
    fig.savefig(save_path + '/inlet_waveform.jpg',dpi=150)
    print('{0} time steps saved, available to visualize in ParaView.'.format(np.unique(np.round(t_pts,3)).shape[0]))
    txt = '{0} time steps saved, available to visualize in ParaView.'.format(np.unique(np.round(t_pts,3)).shape[0])
    return txt