# -*- coding: utf-8 -*-
"""
    Plot the results of a Premod simulation
"""

import numpy as np
import matplotlib.pyplot as plt

from premod.io import Results


class PlotResults(object):
    
    # curve type dictionary: key=curve name/type, value=curve label
    curves = {'standard_f': 'Volume fraction vs Time',
              'standard_r': 'Mean radius vs Time',
              'chemistry': 'Composition vs Time'}
    
    def __init__(self, result=None):
        self.result = Results() if result is None else result
        
    def plot(self, figtyp, fig_kw=dict()):
        subplots = {1: 111, 2: 121, 3: 221, 4: 221}
        types = figtyp.split(',')
        n = len(types)
        if n in subplots:
            gca = fig_kw.pop('gca', False) # use the current axes
            subplot = 0 if gca else subplots[n]
            for typ in types:
                func = getattr(self, 'plot_' + typ, None)
                if func is not None:
                    func(subplot=subplot, **fig_kw)
                else:
                    err = 'Figure type unknown: "{}"'.format(figtyp)
                    raise RuntimeError(err)
                if subplot > 0:
                    subplot += 1
            
    def _plot_standard(self, typ, ax1, ax2,both_axes=True):
        """
        Function to generate a 'standard' plot meaning the time evolution
        of some variables either fv (-);rv_avg (m);Nv (m-3) (old version of Premod)
        in new version MG5SI6_B_DP_fv (-);MG5SI6_B_DP_rv_avg (m); ;MG5SI6_B_DP_Nv (m-3)
        """
        if 'Time' not in self.result.history:
            return
        x = self.result.history['Time']
        if self.result.history.IO_version == 1:
            yname = 'fv' if typ == 'f' else 'rv_avg'
            ylabel = r'$f_v \,\, (-)$' if typ == 'f' else r'$\overline{r}_{v} \,\, (m)$'
            leg_label = r'$f_v$' if typ == 'f' else r'$\overline{r}_{v}$'
            y1 = self.result.history[yname]
            y2 = self.result.history['Nv']
            leg_label2 = r'$N_v$'
        elif self.result.history.IO_version == 2:
            # hard-coded phase
            phase = "MG5SI6_B_DP"
            yname = 'fv' if typ == 'f' else 'rv_avg'
            yname = phase + "_" + yname
            ylabel = r'$f_v \,\, (-)$' if typ == 'f' else r'$\overline{r}_{v} \,\, (m)$'
            leg_label = r'$f_v$' if typ == 'f' else r'$\overline{r}_{v}$'
            leg_label = leg_label + " " + phase
            y1 = self.result.history[yname]
            yname2 = phase+"_Nv"
            y2 = self.result.history[yname2]
            leg_label2 = r'$N_v$' + " " + phase
        else:
            raise ValueError("unrecognized version: ",self.result.history.IO_version)
        
        ax1.plot(x, y1, label=leg_label, color='red')
        ax1.grid(True)
        ax1.set_xscale('log')
        ax1.set_xlabel(r'$Time \,\, (s)$', fontsize=14)
        ax1.set_ylabel(ylabel, fontsize=14)
        lines, labels = ax1.get_legend_handles_labels()
        if both_axes:
            ax2.plot(x, y2, label=leg_label2, color='blue')
            ax2.set_ylabel(r'$N_{v} \,\, (m^{-3})$', fontsize=14)
            lines_2, labels_2 = ax2.get_legend_handles_labels()
        
            lines += lines_2
            labels += labels_2
        
        ax1.legend(lines, labels, loc='best')
        
    def _gca(self, twinx, subplot, **fig_kw):
        """ Create or get the axes to do the plot """
        if subplot:
            plt.figure(**fig_kw)
            ax1 = plt.subplot(subplot)
        else:
            ax1 = plt.gca()
        if twinx:
            return ax1, ax1.twinx()
        else:
            return ax1
        
    def plot_standard_f(self, subplot=111,both_axes=True, **fig_kw):
        """ Function to plot in a standard way data from the ppt history file
            generated by Premod
        """
        ax1, ax2 = self._gca(True, subplot, **fig_kw)
        self._plot_standard('f', ax1, ax2,both_axes=both_axes)
        
    def plot_standard_r(self, subplot=111,both_axes=True, **fig_kw):
        """ Function to plot in a standard way data from the ppt history file
            generated by Premod
        """
        ax1, ax2 = self._gca(True, subplot, **fig_kw)
        self._plot_standard('r', ax1, ax2,both_axes=both_axes)
        
    def plot_chemistry(self, subplot=111, **fig_kw):
        """ Plot composition as function of Time from the ppt history file
            generated by Premod
        """
        self._gca(False, subplot, **fig_kw)
        if 'Time' in self.result.history:
            x = self.result.history['Time']
            if self.result.history.IO_version == 1:
                label_chem = "X_"
            elif self.result.history.IO_version == 2:
                label_chem = "XSS_"
            else:
                raise ValueError("unrecognized version: ",self.result.history.IO_version)
            for key, val in self.result.history.items():
                if label_chem in key:
                    label = key.split(label_chem)[1]
                    if label != 'Al':
                        plt.plot(x, val, label=r'$'+label+'$')
            plt.xscale('log')
            plt.xlabel(r'$Time \,\, (s)$', fontsize=14)
            plt.ylabel(r'$Composition \,\, (at)$', fontsize=14)
            plt.legend(loc='best')
            plt.grid(True)

def plot_sizedistribution(df, subplot=111, **fig_kw):
    """ Plot size distribution from the microstructure file
        generated by Premod
    """
    ax1 = plt.gca()
    if 'rpart[m]' in df.columns:
        x = df['rpart[m]']
        y = df['Nd[1/m3]']
        plt.plot(x, y, label='Nd[1/m3]')
        #plt.xscale('log')
        plt.xlabel('Particle radius (m)', fontsize=14)
        plt.ylabel('Particle density (1/m3)', fontsize=14)
        plt.legend(loc='best')
        plt.grid(True)

    
        