#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
***********************************************************************************
                           tutorial_opencs_dae_7_cv.py
                DAE Tools: pyOpenCS module, www.daetools.com
                Copyright (C) Dragan Nikolic
***********************************************************************************
DAE Tools is free software; you can redistribute it and/or modify it under the
terms of the GNU General Public License version 3 as published by the Free Software
Foundation. DAE Tools is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
PARTICULAR PURPOSE. See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along with the
DAE Tools software; if not, see <http://www.gnu.org/licenses/>.
************************************************************************************
"""
__doc__ = """
Code verification using the Method of Manufactured Solutions.

This problem and its solution in COMSOL Multiphysics <https://www.comsol.com> software
is described in the COMSOL blog: Verify Simulations with the Method of Manufactured Solutions (2015),
https://www.comsol.com/blogs/verify-simulations-with-the-method-of-manufactured-solutions.

Here, a 1D transient heat conduction problem in a bar of length L is solved using the
Centered Finite Difference method:
   dT/dt - k/(rho*cp) * d2T/dx2 = 0, x in [0,L]
Dirichlet boundary conditions are applied at both bar ends (500 K).
Initial conditions are 500 K.

The manufactured solution is given by function u(x):
   u(x) = 500 + (x/L) * (x/L - 1) * (t/tau)

Note:
 csSimulate cannot be called multiple times from python (multiple MPI init and finalize are not allowed).
Therefore, the simulation are performed using the csSimulator binary
(it has to be in the PATH for this to work).
"""

import os, sys, json, itertools, numpy, pandas
from daetools.solvers.opencs import pyOpenCS
from daetools.solvers.opencs import csModelBuilder_t, csNumber_t, csSimulate
from daetools.solvers.opencs import csGraphPartitioner_t, createGraphPartitioner_2D_Npde
from daetools.examples.tutorial_opencs_aux import compareResults
import matplotlib.pyplot as plt

class HeatConductionBar_1D:
    def __init__(self, Nx_):
        self.Nx = Nx_ + 1
        Nx = self.Nx

        self.L = 100.0 # m

        self.x0 = 0.0
        self.x1 = self.L
        self.dx = (self.x1-self.x0) / (Nx-1)

        self.x_domain = []
        for x in range(self.Nx):
            self.x_domain.append(self.x0 + x * self.dx)

        self.T_start_index = 0*Nx
        self.u_start_index = 1*Nx

        self.Nequations = 2*Nx

    def GetInitialConditions(self):
        # Use numpy array so that setting T_0 and u_0 changes the original values
        uv0 = numpy.zeros(self.Nequations)
        T_0 = uv0[self.T_start_index : self.u_start_index]
        u_0 = uv0[self.u_start_index : self.Nequations]

        for ix in range(self.Nx):
            xp = self.x_domain[ix]

            T_0[ix] = 500.0
            u_0[ix] = 500.0

        return uv0.tolist()

    def GetVariableNames(self):
        return ['T(%d)'%x for x in range(self.Nx)] + ['u(%d)'%x for x in range(self.Nx)]

    def CreateEquations(self, y, dydt, currentTime):
        # y is a list of csNumber_t objects representing model variables
        # dydt is a list of csNumber_t objects representing time derivatives of model variables
        T_values    = y   [self.T_start_index : self.u_start_index]
        u_values    = y   [self.u_start_index : self.Nequations]
        dTdt_values = dydt[self.T_start_index : self.u_start_index]
        
        Nx = self.Nx
        dx = self.dx

        #Ac    =    0.1  # m**2
        rho   = 2700.0  # kg/m**3
        cp    =  900.0  # J/(kg*K)
        kappa =  238.0  # W/(m*K)
        tau   = 3600.0  # seconds
        L     =  self.L # m
        t     = currentTime
        # Thermal diffusivity (m**2/s)
        alpha = kappa/(rho * cp)

        # Manufactured solution
        # Note: x is not an index here, but length in meters
        u       = lambda x: 500.0 + (x/L) * (x/L - 1) * (t/tau)
        du_dt   = lambda x: (x/L) * (x/L - 1) * (1/tau)
        d2u_dx2 = lambda x: (2/(L**2)) * (t/tau)
        Q       = lambda x: du_dt(x) - alpha * d2u_dx2(x)

        # Variables
        T       = lambda x:    T_values[x]
        um      = lambda x:    u_values[x]
        dT_dt   = lambda x: dTdt_values[x]

        def d2T_dx2(x):
            # Interior points
            Ti1 = T(x+1)
            Ti  = T(x)
            Ti2 = T(x-1)
            return (Ti1 - 2*Ti + Ti2) / (dx*dx)

        eq = 0
        equations = [None] * self.Nequations
        # Component T:
        for x in range(Nx):
            if(x == 0):       # BCs: Dirichlet
                equations[eq] = T(x) - 500.0
            elif(x == Nx-1):  # BCs: Dirichlet
                equations[eq] = T(x) - 500.0
            else:             # Interior points
                xp = self.x_domain[x]
                equations[eq] = dT_dt(x) - alpha * d2T_dx2(x) - Q(xp)
            eq += 1
        
        # Manuf. solution u:
        for x in range(Nx):
            xp = self.x_domain[x]
            equations[eq] = um(x) - u(xp)
            eq += 1
        
        return equations

def loadOpenCSData(Nx, resultsDirectory = 'tutorial_opencs_dae_7_cv'):
    csvFilename = os.path.join(resultsDirectory, '%d' % Nx, 'results-0.csv')
    csv_filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), csvFilename)
    df = pandas.read_csv(csv_filepath, sep=';', header=2, skiprows=None, quotechar='"', skipinitialspace=True, dtype=float)

    Nx = Nx+1

    print('Loading %d data...' % (Nx-1))

    T_start = 1 # The first value is Time (skip it)
    T_end   = T_start + Nx

    u_start = T_end
    u_end   = u_start + Nx

    data = df.values[-1] # take the last time point
    T  =  data[T_start:T_end]
    um =  data[u_start:u_end]

    return T, um

def simulate(Nx):
    # Instantiate the model being simulated.
    model = HeatConductionBar_1D(Nx)
    
    # 1. Initialise the DAE system with the number of variables and other inputs.
    mb = csModelBuilder_t()
    mb.Initialize_DAE_System(model.Nequations, 0, defaultAbsoluteTolerance = 1e-6)
    
    # 2. Specify the OpenCS model.
    mb.ModelEquations = model.CreateEquations(mb.Variables, mb.TimeDerivatives, mb.Time)
    # Set initial conditions.
    mb.VariableValues = model.GetInitialConditions()
    # Set variable names.
    mb.VariableNames  = model.GetVariableNames()
    
    # 3. Generate a model for single CPU simulations.    
    # Set simulation options (specified as a string in JSON format).
    options = mb.SimulationOptions
    options['Simulation']['OutputDirectory']             = '%d' % Nx
    options['Simulation']['TimeHorizon']                 = 20*3600.0
    options['Simulation']['ReportingInterval']           = 3600.0
    options['Solver']['Parameters']['RelativeTolerance'] = 1e-5
    
    options['LinearSolver'] = {
        "Library":   "Trilinos",
        "Name":      "Amesos_Klu",
        "PrintInfo": False,
        "Parameters": {}
    }
    mb.SimulationOptions = options
    
    # Partition the system to create the OpenCS model for a single CPU simulation.
    # In this case (Npe = 1) the graph partitioner is not required.
    Npe = 1
    graphPartitioner = None
    inputFilesDirectory = 'tutorial_opencs_dae_7_cv'
    cs_models = mb.PartitionSystem(Npe, graphPartitioner)
    csModelBuilder_t.ExportModels(cs_models, inputFilesDirectory, mb.SimulationOptions)
    print("Single CPU OpenCS model generated successfully!")

    # 4. Run simulation using the exported model from the specified directory.
    # Cannot call csSimulate multiple times from python (multiple MPI init and finalize are not allowed).
    # Therefore, simulate it using the csSimulator binary (it has to be in the PATH for this to work).
    #csSimulate(inputFilesDirectory)
    os.system('csSimulator %s' % inputFilesDirectory)

    x_coords = numpy.linspace(0, model.L, Nx+1)
    T, um = loadOpenCSData(Nx)
    return  x_coords, T, um

def run(**kwargs):
    Nx1 = 5
    Nx2 = 20
    L = 100.0
    h1 = L / Nx1
    h2 = L / Nx2
    xcoords1, T1, u1 = simulate(Nx1)
    xcoords2, T2, u2 = simulate(Nx2)

    # The normalised global errors
    E1 = numpy.sqrt((1.0/Nx1) * numpy.sum((T1-u1)**2))
    E2 = numpy.sqrt((1.0/Nx2) * numpy.sum((T2-u2)**2))

    # Order of accuracy
    p = numpy.log(E1/E2) / numpy.log(h1/h2)
    C = E1 / h1**p

    print('\n\nOrder of Accuracy:')
    print('||E(h)|| is proportional to: C * (h**p)')
    print('||E(h1)|| = %e, ||E(h2)|| = %e' % (E1, E2))
    print('C = %e' % C)
    print('Order of accuracy (p) = %.2f' % p)

    fontsize = 14
    fontsize_legend = 11
    fig = plt.figure(figsize=(10,4), facecolor='white')
    #title = 'Plots for coarse and fine grids (Order of accuracy = %.2f) (cv_5)' % (p)
    #fig.canvas.setWindowTitle(title)

    ax = plt.subplot(121)
    plt.plot(xcoords1, T1, 'rs', label='T (CFDM)')
    plt.plot(xcoords1, u1, 'b-', label='u (manufactured)')
    plt.xlabel('x, m', fontsize=fontsize)
    plt.ylabel('Temperature, K', fontsize=fontsize)
    plt.legend(fontsize=fontsize_legend)
    plt.xlim((0, 100))

    ax = plt.subplot(122)
    plt.plot(xcoords2, T2, 'rs', label='T (CFDM)')
    plt.plot(xcoords2, u2, 'b-', label='u (manufactured)')
    plt.xlabel('x, m', fontsize=fontsize)
    plt.ylabel('Temperature, K', fontsize=fontsize)
    plt.legend(fontsize=fontsize_legend)
    plt.xlim((0, 100))

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    run()