# coding: utf-8
"""
Morad Ben Tayeb
CEA Saclay - DM2S/STMF/LGLS
Mars 2021 - Stage 6 mois

The main class is SpatialPostProcessing. Using this package allows the following:

- extract 2D profiles from lata/med files using visit
- plot the 2D spatial profiles
- plot the spatial autocorrelations & calculate the Taylor micro-scale
- plot the spatial energy spectrum using the Welch method

"""
import os
import numpy as np
from trustutils.visitutils import export_lata_base
from .trust_statistical_tools_base import StatisticalPostProcessing


class SpatialPostProcessing(StatisticalPostProcessing):
    """

    Spatial post processing to analyze lata files.
    Mainly, we give the possibility for the user to extract profiles from visit Query.
    All required is to precise the starting and ending point of the segment.
    Once defined, visit will be launched automatically (in terminal) and a .Curve file will be created.
    The idea afterwards is to read these curve files and save them as numpy arrays to plot them, compute
    correlations, spectra, etc. 

    """


    def __init__(self):
        self.temp_folder = "CurveFiles/"
        self.L_NAME = []
        pass

    def __del__(self):
        """

        Removing temporary files, generated by visit. Executed when the object is destroyed (at the end of the script for instance)

        """
        if os.path.exists("tmp.py"):
            os.remove("tmp.py")

        if os.path.exists("visitlog.py"):
            os.remove("visitlog.py")

        if os.path.exists("tmp.out"):
            os.remove("tmp.out")

    def _getFile(self, name, ext=".curve"):
        """

        Return the name of the curve file

        Parameters
        ----------
        name : str 
            name of the saved file 

        return
        ------
        The curvefile name (egg. "name" returns CurveFiles/name.curve)

        """
        return self.temp_folder + name + ext

    def plot(self, x1, y1, label1="", linestyle1="-", linewidth1=0.5, **kwargs):
        """

        Method to plot one are multiple data. Calls the StatisticalPostProcessing plot.

        Parameters
        ----------
        x : array

        y : array

        label : str

        linewidth : float
            To plot multiple data, you can use keywords x1 = ..., y1 = ..., ...

        Returns
        -------
        None

        """
        _save_path = kwargs.get("save_path")

        if _save_path == None:
            _save_path = "profileResults/"

        kwargs["save_path"] = _save_path
        StatisticalPostProcessing._plot(self, x1, y1, label1=label1, linestyle1=linestyle1, linewidth1=linewidth1, **kwargs)

    def extract_profile(self, lata, field, start_point, end_point, name, frame=-1):
        """

        Method to extract profile from by launching visit and save the result as .curve file. 
        By default frame = -1 that means the last frame of the simulation.

        Parameters
        ----------
        lata : str
            Path to the lata file.
        field : str
            Name of the field. Example: "VITESSE_X_ELEM_dom" 
        start_point : tuple
            Starting point of the segment. Example: (0.,0.,0.)
        end_point : tuple 
            Ending point of the segment. Example: (0.1,0.,0.)
        name : str 
            Name of the saved . curve file
        frame : int 
            Selected frame, frame = -1 for the last frame, optional.  

        Returns
        -------
        None

        """
        StatisticalPostProcessing._checkPath(self, lata)
        StatisticalPostProcessing._checkName(self, name)
        if not os.path.exists(self.temp_folder):
            os.mkdir(self.temp_folder)

        StatisticalPostProcessing._COUNT_CURVE_FILES += 1
        elb = export_lata_base(lata, "Pseudocolor", field, self.temp_folder + name, frame=frame)
        SpatialPostProcessing.L_NAME += [name]
        # elb.lineout(start_point,end_point)
        elb.query(start_point, end_point)
        elb.save()
        return ()

    def getDimensions(self, lata, field, frame=-1):
        """

        Get dimensions, [Nx, Ny, Nz].

        """
        elb = export_lata_base(lata, "Pseudocolor", field, None, frame=frame)
        return elb.getDimensions()

    def getFrames(self, lata, field):
        """
        
        Get total number of frames 

        Parameters
        ----------
        name : str 
            Name of the saved . curve file

        Returns
        -------
        N : int
            Number of frames

        """
        elb = export_lata_base(lata, "Pseudocolor", field, self.temp_folder + "frames")
        return elb.getFrames()

    def getSpatialValues(self, name):
        """

        Return the spatial values.

        Parameters
        ----------
        name : str
            Name of the .curve file gived by the user in extract_profile method.

        Returns
        -------
        X : array
            Distance.
        Y : array
            Field.

        """
        X, Y = self._getValuesFromCurve(name)
        return (X, Y)

    def getField(self, name, format="IJK", dimX=-1, dimY=-1, dimZ=-1):
        """

        Return the fields.

        Parameters
        ----------
        name : str
            Name of the .okc file gived by the user in extract_field method or already stored.

        Returns
        -------
        field : array
            Return the field in the format asked. 
        X : array
            X coordinates. If format = "XYZ".
        Y : array
            Y coordinates. If format = "XYZ".
        Z : array
            Z coordinates. If format = "XYZ".


        """
        X, Y, Z, field = self._getFieldsFromOKC(name)

        if format == "XYZ":
            return (field, X, Y, Z)

        if format == "IJK":
            if dimX <= 0 or dimY <= 0 or dimZ <= 0:
                raise Exception("Wrong dimensions input.")

            return field.reshape((dimX, dimY, dimZ)).T

    def spatialAutoCorrelation(self, X, Y, max_lag=None, fit=False):
        """

        Computes spatial signal autocorrelation.

        Parameters
        ----------
        X : array
            Distances.
        Y : array 
            Field.
        max_lag : float, Optional
            Maximum lag value
        fit : boolean
            Set to ``True`` to perform a parbola fit on autocorrelation 

        Returns
        -------
        dist : array
            Distances.
        autocorr : array
            Array of signal autocorrelation.
        Xfit : array
            Result of parabola fit for X-axis, only if ``fit = True``
        Yfit : array
            The parabola fit, only if ``fit = True``

        """
        dist, autocorr = StatisticalPostProcessing._computeAutocorrelation(self, X, Y, max_lag)
        res = [dist, autocorr]
        if fit:
            Xfit, Yfit = StatisticalPostProcessing._fitParabola(self, res[0], res[1])
            res = res + [Xfit, Yfit]
        return res

    def spatialEnergySpectrum(self, X, Y):
        """

        Computes spatial energy spectrum with Welch method.

        Parameters
        ----------
        X : array
            Distances.
        Y : array
            Field.

        Returns
        -------
        lambd : array 
            Wave number.
        power : array 
            Energy.
        
        """
        lambd, power = StatisticalPostProcessing._computeWelch(self, X, Y)
        return (lambd, power)

    def extract_field(self, lata, field, name, frame=-1):
        """

        Method to extract field from by launching visit and save the result as .okc file. 
        By default, frame = -1 that means we choose the last frame of the simulation.

        Parameters
        ----------
        lata : str
            Path to the lata file.
        field : str
            Name of the field. Example: "VITESSE_X_ELEM_dom" 
        name : str 
            Name of the saved .okc file
        frame : int 
            Selected frame, frame = -1 for the last frame, optional.  

        Returns
        -------
        None

        """
        StatisticalPostProcessing._checkPath(self, lata)
        StatisticalPostProcessing._checkName(self, name)
        if not os.path.exists(self.temp_folder):
            os.mkdir(self.temp_folder)

        # StatisticalPostProcessing._COUNT_CURVE_FILES += 1
        elb = export_lata_base(lata, "Pseudocolor", field, self.temp_folder + name)
        elb.setFrame(frame)
        SpatialPostProcessing.L_NAME += [name]
        elb.save(type="Xmdv")
        return ()
