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

"""
***********************************************************************************
                           tutorial_opencs_dae_5_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.
It represents re-implementation of DAE Tools code verifications tutorial CV-2 -
the 1-D transient convection-diffusion (Burger's) equation.
This example uses Dirichlet boundary conditions.

References:

1. G. Tryggvason. Method of Manufactured Solutions, Lecture 33: Predictivity-I, 2011.
   `PDF <http://www3.nd.edu/~gtryggva/CFD-Course/2011-Lecture-33.pdf>`_
2. K. Salari and P. Knupp. Code Verification by the Method of Manufactured Solutions.
   SAND2000 – 1444 (2000).
   `doi:10.2172/759450 <https://doi.org/10.2172/759450>`_
3. P.J. Roache. Fundamentals of Verification and Validation. Hermosa, 2009.
   `ISBN-10:0913478121 <http://www.isbnsearch.org/isbn/0913478121>`_

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 Burgers_1D:
    def __init__(self, Nx_):
        self.Nx = Nx_ + 1
        Nx = self.Nx

        self.A = 1.0
        self.C = 1.0
        self.D = 0.05

        self.x0 = 0.0
        self.x1 = 2*numpy.pi
        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.f_start_index = 0*Nx
        self.q_start_index = 1*Nx

        self.Nequations = 2*Nx

    def GetInitialConditions(self):
        # Use numpy array so that setting f_0 and q_0 changes the original values
        uv0 = numpy.zeros(self.Nequations)
        f_0 = uv0[self.f_start_index : self.q_start_index]
        q_0 = uv0[self.q_start_index : self.Nequations]

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

            f_0[ix] = self.A + numpy.sin(xp)
            q_0[ix] = self.A + numpy.sin(xp)

        return uv0.tolist()

    def GetVariableNames(self):
        return ['f(%d)'%x for x in range(self.Nx)] + ['q(%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
        f_values    = y   [self.f_start_index : self.q_start_index]
        q_values    = y   [self.q_start_index : self.Nequations]
        dfdt_values = dydt[self.f_start_index : self.q_start_index]
        
        dx = self.dx
        Nx = self.Nx
        xp = self.x_domain

        A = self.A
        C = self.C
        D = self.D
        t = currentTime

        # Math. functions
        sin    = pyOpenCS.sin
        cos    = pyOpenCS.cos

        # Variables
        f       = lambda x:    f_values[x]
        qm      = lambda x:    q_values[x]
        df_dt   = lambda x: dfdt_values[x]

        def df_dx(x):
            # Interior points
            f1 = f(x+1)
            f2 = f(x-1)
            return (f1 - f2) / (2*dx)
        def d2f_dx2(x):
            # Interior points
            fi1 = f(x+1)
            fi  = f(x)
            fi2 = f(x-1)
            return (fi1 - 2*fi + fi2) / (dx*dx)

        q       = lambda x: A + sin(xp[x] + C*t)
        dq_dt   = lambda x: C * cos(xp[x] + C*t)
        dq_dx   = lambda x: cos(xp[x] + C*t)
        d2q_dx2 = lambda x: -sin(xp[x] + C*t)
        g       = lambda x: dq_dt(x) + q(x) * dq_dx(x) - D * d2q_dx2(x)

        eq = 0
        equations = [None] * self.Nequations
        # Component u:
        for x in range(Nx):
            if(x == 0):       # BCs: Dirichlet
                equations[eq] = f(x) - q(x)
            elif(x == Nx-1):  # BCs: Dirichlet
                equations[eq] = f(x) - q(x)
            else:
                # Interior points
                equations[eq] = df_dt(x) + f(x) * df_dx(x) - D * d2f_dx2(x) - g(x)
            eq += 1
        
        # Component v:
        for x in range(Nx):
            equations[eq] = qm(x) - q(x)
            eq += 1
        
        return equations

def loadOpenCSData(Nx, resultsDirectory = 'tutorial_opencs_dae_5_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))

    f_start = 1 # The first value is Time (skip it)
    f_end   = f_start + Nx

    q_start = f_end
    q_end   = q_start + Nx

    data = df.values[-1] # take the last time point
    f  =  data[f_start:f_end]
    qm =  data[q_start:q_end]

    return f, qm

def simulate(Nx):
    # Instantiate the model being simulated.
    model = Burgers_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']                 = 1.0
    options['Simulation']['ReportingInterval']           = 0.05
    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_5_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)

    f, q = loadOpenCSData(Nx)
    return  f, q

def run(**kwargs):
    Nxs = numpy.array([60, 90, 120, 150])
    n = len(Nxs)
    L = 2*numpy.pi
    hs = L / Nxs
    E = numpy.zeros(n)
    C = numpy.zeros(n)
    p = numpy.zeros(n)
    E2 = numpy.zeros(n)

    # The normalised global errors
    for i,Nx in enumerate(Nxs):
        numerical_sol, manufactured_sol = simulate(int(Nx))
        E[i] = numpy.sqrt((1.0/Nx) * numpy.sum((numerical_sol-manufactured_sol)**2))

    # Order of accuracy
    for i,Nx in enumerate(Nxs):
        p[i] = numpy.log(E[i]/E[i-1]) / numpy.log(hs[i]/hs[i-1])
        C[i] = E[i] / hs[i]**p[i]

    C2 = 0.09 # constant for the second order slope line (to get close to the actual line)
    E2 = C2 * hs**2 # E for the second order slope

    fontsize = 14
    fontsize_legend = 11
    fig = plt.figure(figsize=(10,4), facecolor='white')
    #fig.canvas.setWindowTitle('The Normalised global errors and the Orders of accuracy (Nelems = %s) (cv_2)' % Nxs.tolist())

    ax = plt.subplot(121)
    plt.figure(1, facecolor='white')
    plt.loglog(hs, E,  'ro', label='E(h)')
    plt.loglog(hs, E2, 'b-', label='2nd order slope')
    plt.xlabel('h', fontsize=fontsize)
    plt.ylabel('||E||', fontsize=fontsize)
    plt.legend(fontsize=fontsize_legend)
    plt.xlim((0.04, 0.11))

    ax = plt.subplot(122)
    plt.figure(1, facecolor='white')
    plt.semilogx(hs[1:], p[1:],  'rs-', label='Order of Accuracy (p)')
    plt.xlabel('h', fontsize=fontsize)
    plt.ylabel('p', fontsize=fontsize)
    plt.legend(fontsize=fontsize_legend)
    plt.xlim((0.04, 0.075))
    plt.ylim((2.0, 2.03))

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    run()