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

"""
***********************************************************************************
                           tutorial_opencs_dae_9.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__ = """
Reimplementation of IDAS idasHeat2D_kry example.
The DAE system solved is a spatial discretization of the PDE::

   du/dt = d^2u/dx^2 + d^2u/dy^2

on the unit square.
The boundary condition is u = 0 on all edges.
Initial conditions are given by u = 16 x (1 - x) y (1 - y).
The PDE is treated with central differences on a uniform M x M grid.
The values of u at the interior points satisfy ODEs, and
equations u = 0 at the boundaries are appended, to form a DAE
system of size N = M^2.
The grid size is 10 x 10.
"""

import os, sys, itertools, numpy, pandas, time
from daetools.solvers.opencs import csModelBuilder_t, csSimulate
import matplotlib.pyplot as plt
from matplotlib import animation

class HeatEquation_2D:
    def __init__(self, Nx, Ny):
        self.Nx = Nx
        self.Ny = Ny
        
        self.x0 = 0.0
        self.x1 = 1.0
        self.y0 = 0.0
        self.y1 = 1.0
        self.dx = (self.x1-self.x0) / (Nx-1)
        self.dy = (self.y1-self.y0) / (Ny-1)

        self.x_domain = []
        self.y_domain = []
        for x in range(self.Nx):
            self.x_domain.append(self.x0 + x * self.dx)
        for y in range(self.Ny):
            self.y_domain.append(self.y0 + y * self.dy)

        self.u_start_index  = 0*Nx*Ny

        self.Nequations = Nx*Ny

    def GetInitialConditions(self):
        # Use numpy array so that setting u_0 and v_0 changes the original values
        u0  = numpy.zeros(self.Nequations)

        for ix in range(self.Nx):
            for iy in range(self.Ny):
                index = self.GetIndex(ix,iy)

                x = self.x_domain[ix]
                y = self.y_domain[iy]
                #u0[index] = 300; #numpy.sin(2*numpy.pi*x) * numpy.sin(2*numpy.pi*y)
                u0[index] = 16 * x * (1 - x) * y * (1 - y)

        return u0.tolist()

    def GetVariableNames(self):
        x_y_inds = [(x,y) for x,y in itertools.product(range(self.Nx), range(self.Ny))]
        return ['u(%d,%d)'%(x,y) for x,y in x_y_inds]

    def CreateEquations(self, y, dydt, t):
        # 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
        u_values     = y   [self.u_start_index  : self.Nequations]
        dudt_values  = dydt[self.u_start_index  : self.Nequations]
        
        dx = self.dx
        dy = self.dy
        Nx = self.Nx
        Ny = self.Ny

        def u(x, y):
            index = self.GetIndex(x, y)
            return u_values[index]

        def du_dt(x, y):
            index = self.GetIndex(x, y)
            return dudt_values[index]

        # Second order partial derivative per x.
        def d2u_dx2(x, y):
            if(x == 0 or x == Nx-1):
                raise RuntimeError("d2u_dx2 called at the boundary")

            ui1 = u(x+1, y)
            ui  = u(x,   y)
            ui2 = u(x-1, y)
            return (ui1 - 2*ui + ui2) / (dx*dx)

        # Second order partial derivative per y.
        def d2u_dy2(x, y):
            if(y == 0 or y == Ny-1):
                raise RuntimeError("d2u_dy2 called at the boundary")

            ui1 = u(x, y+1)
            ui  = u(x,   y)
            ui2 = u(x, y-1)
            return (ui1 - 2*ui + ui2) / (dy*dy)

        eq = 0
        equations = [None] * self.Nequations
        # Component u:
        for x in range(Nx):
            for y in range(Ny):
                if(x == 0):       # Left BC: Dirichlet BCs
                    xp = self.x_domain[x]
                    equations[eq] = u(x,y) - 0
                elif(x == Nx-1):  # Right BC: Dirichlet BCs
                    equations[eq] = u(x,y) - 0
                elif(y == 0):     # Bottom BC: Dirichlet BCs
                    equations[eq] = u(x,y) - 0
                elif(y == Ny-1):  # Top BC: Dirichlet BCs
                    equations[eq] = u(x,y) - 0
                else:
                    # Interior points
                    equations[eq] = du_dt(x,y) - (d2u_dx2(x,y) + d2u_dy2(x,y))
                eq += 1

        return equations
    
    def GetIndex(self, x, y):
        if x < 0 or x >= self.Nx:
            raise RuntimeError("Invalid x index")
        if y < 0 or y >= self.Ny:
            raise RuntimeError("Invalid y index")
        return self.Ny*x + y
    
def run(**kwargs):
    inputFilesDirectory = kwargs.get('inputFilesDirectory', os.path.splitext(os.path.basename(__file__))[0])
    Nx         = kwargs.get('Nx', 50)
    Ny         = kwargs.get('Ny', 50)

    # Instantiate the model being simulated.
    model = HeatEquation_2D(Nx, Ny)
    
    # 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-5)
    
    # 2. Specify the OpenCS model.
    # Create and set model equations using the provided time/variable/timeDerivative/dof objects.
    # The DAE system is defined as:
    #     F(x',x,y,t) = 0
    # where x' are derivatives of state variables, x are state variables,
    # y are fixed variables (degrees of freedom) and t is the current simulation time.
    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']             = 'results'
    options['Simulation']['TimeHorizon']                 = 10.24
    options['Simulation']['ReportingInterval']           =  0.08
    options['Solver']['Parameters']['RelativeTolerance'] =  1e-5
    
    # ILU options for Ncpu = 1: k = 3, rho = 1.0, alpha = 1e-1, w = 0.5
    options['LinearSolver']['Preconditioner']['Parameters']['fact: level-of-fill']      =    3
    options['LinearSolver']['Preconditioner']['Parameters']['fact: relax value']        =  0.5
    options['LinearSolver']['Preconditioner']['Parameters']['fact: absolute threshold'] = 1e-1
    options['LinearSolver']['Preconditioner']['Parameters']['fact: relative threshold'] =  1.0
    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
    cs_models = mb.PartitionSystem(Npe, graphPartitioner)
    csModelBuilder_t.ExportModels(cs_models, inputFilesDirectory, mb.SimulationOptions)
    print("Single CPU OpenCS model generated successfully!")

    # 5. Run simulation using the exported model from the specified directory.
    csSimulate(inputFilesDirectory)
    
    try:
        csvFilename = os.path.join(inputFilesDirectory, 'results', '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)

        # The data layout in memory is as (for each time point):
        #   [Time][u data]
        #   [  1 ][ N*Ny ]

        # Offset the data by one point (time).
        c_start   = 1
        c_end     = c_start + Nx*Ny

        Nframes = len(df.values)

        # Plot setup and the first frame
        data = df.values[0]
        t = data[0]
        c = data[c_start : c_end]
        c = numpy.array(c).reshape(Nx,Ny)
        fig = plt.figure()
        im = plt.imshow(c, cmap = 'viridis')
        plt.colorbar()
        plt.title('time %.2f s' % t)

        def animate(i):
            data = df.values[i]
            t = data[0]
            c = data[c_start : c_end]
            c = numpy.array(c).reshape(Nx,Ny)
            im.set_data(c)
            plt.title('time %.2f s' % t)
            time.sleep(0.5)
            return im

        plt.__anim__ = animation.FuncAnimation(fig, animate, frames=Nframes, repeat = False)
        plt.show()
    except Exception as e:
        print(str(e))

if __name__ == "__main__":
    if len(sys.argv) == 1:
        Nx = 50
        Ny = 50
    elif len(sys.argv) == 3:
        Nx = int(sys.argv[1])
        Ny = int(sys.argv[2])
    else:
        print('Usage: python tutorial_opencs_dae_9.py Nx Ny')
        sys.exit()
        
    inputFilesDirectory = 'tutorial_opencs_dae_9'
    run(Nx = Nx,
        Ny = Ny,
        inputFilesDirectory = inputFilesDirectory)