import re

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from pathlib import Path

from matplotlib.patches import Patch
from matplotlib.lines import Line2D

from pyfluids import Fluid, FluidsList, Input

plt.rcParams.update({'font.size': 22})
plt.rcParams.update({'mathtext.fontset':'cm'})
plt.rcParams.update({'font.family':'Times New Roman'})

def make_dict(paths, Dh, Pr, density, nu, Cp, k, H, A, h_area, Twall, solver):

    def find_nearest(array, value):
        array = np.asarray(array)
        idx = (np.abs(array - value)).argmin()
        return array[idx]

    plot_dict = {}

    for i, path in enumerate(paths):
        case = path.stem
        plot_dict[case] = {'periodic':{}, 'analytical':{}}
        # analytical = TemperatureField(1.0, 0.5, 0.7, float(case[2::]), nu, True, 8)    
        df = pd.read_csv(path)
        x_coords = df['Points_0'].unique()
        min_x = np.min(x_coords)
        max_x = np.max(x_coords)
        plot_dict[case]['periodic']['L'] = max_x-min_x
        mid_x = find_nearest(x_coords, plot_dict[case]['periodic']['L']/2)
        
        inlet_df = df[df['Points_0']==min_x]
        outlet_df = df[df['Points_0']==max_x]
        midway_df = df[df['Points_0']==mid_x]

        plot_dict[case]['periodic']['x'] = df['Points_0'].values-min_x
        plot_dict[case]['periodic']['y'] = df['Points_1'].values   
        plot_dict[case]['periodic']['yI'] = inlet_df['Points_1'].values
        plot_dict[case]['periodic']['yO'] = outlet_df['Points_1'].values
        plot_dict[case]['periodic']['yM'] = midway_df['Points_1'].values

        plot_dict[case]['periodic']['Ux'] = df['UMean_0'].values
        plot_dict[case]['periodic']['Umean'] = np.mean(plot_dict[case]['periodic']['Ux'])
        plot_dict[case]['periodic']['UI'] = inlet_df['UMean_0'].values
        plot_dict[case]['periodic']['UImean'] = np.mean(plot_dict[case]['periodic']['UI'])
        plot_dict[case]['periodic']['UO'] = outlet_df['UMean_0'].values
        plot_dict[case]['periodic']['UM'] = midway_df['UMean_0'].values
        plot_dict[case]['periodic']['UMmean'] = np.mean(plot_dict[case]['periodic']['UM'])
        
        plot_dict[case]['periodic']['m_dot'] = density*A*plot_dict[case]['periodic']['UImean']
        
        plot_dict[case]['periodic']['T'] = df['TMean'].values
        plot_dict[case]['periodic']['TI'] = inlet_df['TMean'].values
        plot_dict[case]['periodic']['TO'] = outlet_df['TMean'].values
        plot_dict[case]['periodic']['TM'] = midway_df['TMean'].values

        #---------------------- h for periodic ------------------------------------------------
        periodic_TI = (np.sum(np.multiply(plot_dict[case]['periodic']['TI'], plot_dict[case]['periodic']['UI']))/
                       np.sum(plot_dict[case]['periodic']['UI']))
        plot_dict[case]['periodic']['TImc'] = periodic_TI
        periodic_TO = (np.sum(np.multiply(plot_dict[case]['periodic']['TO'], plot_dict[case]['periodic']['UO']))/
                       np.sum(plot_dict[case]['periodic']['UO']))
        plot_dict[case]['periodic']['TOmc'] = periodic_TI

        temp_lm = ((Twall - periodic_TO) - (Twall - periodic_TI))/(np.log((Twall - periodic_TO)/(Twall-periodic_TI)))
        
        plot_dict[case]['periodic']['h'] = np.abs(plot_dict[case]['periodic']['m_dot']*Cp*(periodic_TI-periodic_TO)/(h_area*temp_lm))
        plot_dict[case]['periodic']['Nu'] = plot_dict[case]['periodic']['h']*Dh/k

        plot_dict[case]['periodic']['Pdrop'] = np.abs(np.mean(solver[case]['pg'][-1000:-1]))*density
        
        plot_dict[case]['periodic']['f'] = ((plot_dict[case]['periodic']['Pdrop']*Dh)/
                                            (plot_dict[case]['periodic']['L']*2*density*plot_dict[case]['periodic']['Umean']**2))

        plot_dict[case]['periodic']['Re'] = plot_dict[case]['periodic']['Umean']*Dh/nu
        print(plot_dict[case]['periodic']['Umean'], plot_dict[case]['periodic']['Re'])
    return plot_dict

def search_str(file_path):
    time = []
    pg = []
    with open(file_path, 'r') as f:
        for line in f:
            if re.match(r'Time = ', line):
                time.append(float(line.split()[2]))
            elif re.match(r'Pressure gradient', line):
                pg.append(float(line.split()[-1]))
    
    return np.array(time), np.array(pg)


flow_type = 'arc'
subcase = 'steady'
fluid = 'air'

data_directory = Path(r'./').joinpath(flow_type, 'GCI')

cgi_files = sorted(list(data_directory.rglob('*.csv')))
solver_files = sorted(list(data_directory.rglob('log.cyclic*Foam')))
print(len(cgi_files))

save_path = Path(r'./')

# grab the pressure gradient
solver = {}
for file in solver_files:
    case = file.parent.stem
    solver[case] = {}
    solver[case]['time'], solver[case]['pg'] = search_str(file)


if fluid == 'air':
    working_fluid = Fluid(FluidsList.Air).with_state(
        Input.pressure(101325), Input.temperature(35)
    )
else:
    working_fluid = Fluid(FluidsList.Water).with_state(
        Input.pressure(101325), Input.temperature(35)
    )

Twall = 293.15

plot_dict = {}

if flow_type == 'arc':
    L = 1.0
    Hmin = 0.214286
    Hmax = 0.7138
    Havg = (Hmin/2 + Hmax/2)
    W = 0.02

    A = Havg*W
    h_area = 0.0455876
    Dh = Havg
else:
    L = 1.0
    Hmin = 0.214286
    Hmax = 0.7134
    Havg = (Hmin/2 + Hmax/2)
    W = 0.02

    A = Havg*W
    h_area = 0.0455655
    Dh = Havg

H = 1.0
print(Havg)
Pr = working_fluid.prandtl
density = working_fluid.density
mu = working_fluid.dynamic_viscosity
nu = working_fluid.dynamic_viscosity/working_fluid.density 
Cp = working_fluid.specific_heat
alpha = nu/Pr
k = Cp*mu/Pr

plot_dict = make_dict(cgi_files, Dh, Pr, density, nu, Cp, k, H, A, h_area, Twall, solver)

if flow_type == 'arc':
    refinement_ratio = 2.0
else:
    refinement_ratio = 2.0
    
safety_factor = 1.25

grid = {'grid0':0, 'grid1':1.0, 'grid2':np.round(refinement_ratio,2), 'grid3':np.round(refinement_ratio**2,2)}
colors=['C0','C1','C2','C3']
markers=['o','v','^','*']


plt.rcParams.update({'font.size':22})
fig_x = 12
fig_y = 10

asymptotic = []

#-----------------------------------------------------------------
# first quarter through the domain
#-----------------------------------------------------------------
fig1, ax1 = plt.subplots(figsize=(fig_x, fig_y))

f_1 = np.mean(plot_dict['grid1']['periodic']['TI'])
f_2 = np.mean(plot_dict['grid2']['periodic']['TI'])
f_3 = np.mean(plot_dict['grid3']['periodic']['TI'])

p = np.log(np.abs(f_3-f_2)/np.abs(f_2-f_1))/np.log(refinement_ratio)

rp = refinement_ratio**p

f_h0 = f_1 + ((f_1 - f_2)/(rp - 1))

ax1.scatter(grid['grid0'], f_h0, marker=markers[0], color=colors[1])
ax1.scatter(grid['grid1'], f_1, marker=markers[0], color=colors[0])
ax1.scatter(grid['grid2'], f_2, marker=markers[0], color=colors[0])
ax1.scatter(grid['grid3'], f_3, marker=markers[0], color=colors[0])

# add a plot for the CGI
# These are in unitless percentages of the figure size. (0,0 is bottom left)
left, bottom, width, height = [0.27, 0.2, 0.2, 0.2]
axins = fig1.add_axes([left, bottom, width, height])

GCI_23 = safety_factor*np.abs((f_2 - f_3)/f_2)/(rp - 1)*100
GCI_12 = safety_factor*np.abs((f_1 - f_2)/f_1)/(rp - 1)*100
GCI_R1 = safety_factor*np.abs((f_h0 - f_1)/f_h0)/(rp - 1)*100

asymptotic = GCI_23/(rp*GCI_12)
print('order of convergence: %f, rp: %f, asymptotic: %f' % (p, rp, asymptotic))

axins.scatter('$GCI_{2,3}$', GCI_23, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{1,2}$', GCI_12, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{R,1}$', GCI_R1, marker=markers[0], color=colors[0])

# ax2.set_xlabel('GCI')
axins.set_ylabel('Grid Convergence \n Percentage')

ax1.ticklabel_format(useOffset=False)
ax1.set_xlabel('Normalized Grid Spacing')
ax1.set_ylabel('Avg. Inlet Temperature [k]')

fig1.savefig(save_path.joinpath(flow_type + '_cgi_average_Temperature.pdf'), dpi=300, facecolor='w', bbox_inches='tight')

fig2, ax2 = plt.subplots(figsize=(fig_x, fig_y))

f_1 = plot_dict['grid1']['periodic']['Umean']
f_2 = plot_dict['grid2']['periodic']['Umean']
f_3 = plot_dict['grid3']['periodic']['Umean']

print(f_1, f_2, f_3)

p = np.log(np.abs(f_3-f_2)/np.abs(f_2-f_1))/np.log(refinement_ratio)

rp = refinement_ratio**p

f_h0 = f_1 + ((f_1 - f_2)/(rp - 1))

ax2.scatter(grid['grid0'], f_h0*1000, marker=markers[0], color=colors[1])
ax2.scatter(grid['grid1'], f_1*1000, marker=markers[0], color=colors[0])
ax2.scatter(grid['grid2'], f_2*1000, marker=markers[0], color=colors[0])
ax2.scatter(grid['grid3'], f_3*1000, marker=markers[0], color=colors[0])

# add a plot for the CGI
# These are in unitless percentages of the figure size. (0,0 is bottom left)
left, bottom, width, height = [0.27, 0.3, 0.2, 0.2]
axins = fig2.add_axes([left, bottom, width, height])

GCI_23 = safety_factor*np.abs((f_2 - f_3)/f_2)/(rp - 1)*100
GCI_12 = safety_factor*np.abs((f_1 - f_2)/f_1)/(rp - 1)*100
GCI_R1 = safety_factor*np.abs((f_h0 - f_1)/f_h0)/(rp - 1)*100

asymptotic = GCI_23/(rp*GCI_12)
print('order of convergence: %f, rp: %f, asymptotic: %f' % (p, rp, asymptotic))

axins.scatter('$GCI_{2,3}$', GCI_23, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{1,2}$', GCI_12, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{R,1}$', GCI_R1, marker=markers[0], color=colors[0])

# ax2.set_xlabel('GCI')
axins.set_ylabel('Grid Convergence \n Percentage')

ax2.ticklabel_format(useOffset=False)
ax2.set_xlabel('Normalized Grid Spacing')
ax2.set_ylabel('Avg. Inlet $\mathrm{U_x}$ [mm/s]')

fig2.savefig(save_path.joinpath(flow_type + '_cgi_mean_Ux.pdf'), dpi=300, facecolor='w', bbox_inches='tight')

fig2, ax2 = plt.subplots(figsize=(fig_x, fig_y))

asymptotic = []

colors=['C0','C1','C2','C3']
markers=['o','v','^','*']

f_1 = plot_dict['grid1']['periodic']['h']
f_2 = plot_dict['grid2']['periodic']['h']
f_3 = plot_dict['grid3']['periodic']['h']

p = np.log(np.abs(f_3-f_2)/np.abs(f_2-f_1))/np.log(refinement_ratio)

rp = refinement_ratio**p

f_h0 = f_1 + ((f_1 - f_2)/(rp - 1))

ax2.scatter(grid['grid0'], f_h0, marker=markers[0], color=colors[1])
ax2.scatter(grid['grid1'], f_1, marker=markers[0], color=colors[0])
ax2.scatter(grid['grid2'], f_2, marker=markers[0], color=colors[0])
ax2.scatter(grid['grid3'], f_3, marker=markers[0], color=colors[0])

# add a plot for the CGI
# These are in unitless percentages of the figure size. (0,0 is bottom left)
left, bottom, width, height = [0.27, 0.2, 0.2, 0.2]
axins = fig2.add_axes([left, bottom, width, height])

GCI_23 = safety_factor*np.abs((f_2 - f_3)/f_2)/(rp - 1)*100
GCI_12 = safety_factor*np.abs((f_1 - f_2)/f_1)/(rp - 1)*100
GCI_R1 = safety_factor*np.abs((f_h0 - f_1)/f_h0)/(rp - 1)*100

asymptotic = GCI_23/(rp*GCI_12)
print('order of convergence: %f, rp: %f, asymptotic: %f' % (p, rp, asymptotic))

axins.scatter('$GCI_{2,3}$', GCI_23, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{1,2}$', GCI_12, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{R,1}$', GCI_R1, marker=markers[0], color=colors[0])

# ax2.set_xlabel('GCI')
axins.set_ylabel('Grid Convergence \n Percentage')

ax2.set_xlabel('Normalized Grid Spacing')
ax2.set_ylabel('Avg. Heat Transfer Coefficient')

fig2.savefig(save_path.joinpath(flow_type + '_cgi_lhtc.pdf'), dpi=300, facecolor='w', bbox_inches='tight')

fig3, ax3 = plt.subplots(figsize=(fig_x, fig_y))

asymptotic = []

colors=['C0','C1','C2','C3']
markers=['o','v','^','*']

f_1 = plot_dict['grid1']['periodic']['f']
f_2 = plot_dict['grid2']['periodic']['f']
f_3 = plot_dict['grid3']['periodic']['f']

p = np.log(np.abs(f_3-f_2)/np.abs(f_2-f_1))/np.log(refinement_ratio)

rp = refinement_ratio**p

f_h0 = f_1 + ((f_1 - f_2)/(rp - 1))

print(f_h0)
ax3.scatter(grid['grid0'], f_h0, marker=markers[0], color=colors[1])
ax3.scatter(grid['grid1'], f_1, marker=markers[0], color=colors[0])
ax3.scatter(grid['grid2'], f_2, marker=markers[0], color=colors[0])
ax3.scatter(grid['grid3'], f_3, marker=markers[0], color=colors[0])

# add a plot for the CGI
# These are in unitless percentages of the figure size. (0,0 is bottom left)
left, bottom, width, height = [0.27, 0.6, 0.2, 0.2]
axins = fig3.add_axes([left, bottom, width, height])

GCI_23 = safety_factor*np.abs((f_2 - f_3)/f_2)/(rp - 1)*100
GCI_12 = safety_factor*np.abs((f_1 - f_2)/f_1)/(rp - 1)*100
GCI_R1 = safety_factor*np.abs((f_h0 - f_1)/f_h0)/(rp - 1)*100

asymptotic = GCI_23/(rp*GCI_12)
print('order of convergence: %f, rp: %f, asymptotic: %f' % (p, rp, asymptotic))

axins.scatter('$GCI_{2,3}$', GCI_23, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{1,2}$', GCI_12, marker=markers[0], color=colors[0])
axins.scatter('$GCI_{R,1}$', GCI_R1, marker=markers[0], color=colors[0])

# ax2.set_xlabel('GCI')
axins.set_ylabel('Grid Convergence \n Percentage')

ax3.set_xlabel('Normalized Grid Spacing')
ax3.set_ylabel('f')

fig3.savefig(save_path.joinpath(flow_type + '_cgi_friction_factor.pdf'), dpi=300, facecolor='w', bbox_inches='tight')