__author__ = 'Miguel Molero'

import os
import sys
import numpy as np
# import matplotlib.pyplot as plt
from  scipy.misc import imsave
from scipy.io import savemat

import math
import yaml

class SnapShots:
    def __init__(self, Enable=False, Extension='.png', Step=200, Filename=None, File_path = None, dB=60, Color=0, 
                       Field = 0,
                       Material = 0,
                       enableFields=False,
                       enableSignals=False,
                       enableImages=False,
                       enableMaterial=False,
                       enableNumPy=False,
                       enableCsv=False,
                       enableVolume=False,
                       enableView=False,
                       enableSummary=False,
                       sensorShape = [],
                       sensorPlacement = [],
                       sensorSize = 0):
        if Material is 0:
          Material = "Im"
        if type(Field) == str or type(Field) == unicode:
          Field = ["Vx","Vy","[Vx,Vy]","Txx","Txy","Tyy","[Txx:Txy:Tyy]","Dx","Dy","[Dx,Dy]","SV"].index(Field)

        self.IsEnable = Enable
        self.Extension = Extension
        self.Step = Step
        self.Filename = Filename
        self.File_path = File_path
        self.DB = dB
        self.Color = Color
        self.Field = Field
        self.Material = Material
        self.enableFields   = enableFields
        self.enableSignals  = enableSignals
        self.enableImages   = enableImages
        self.enableMaterial = enableMaterial
        self.enableNumPy    = enableNumPy
        self.enableCsv      = enableCsv
        self.enableVolume   = enableVolume
        self.enableView     = enableView
        self.enableSummary  = enableSummary
        self.sensorShape    = sensorShape
        self.sensorPlacement = sensorPlacement
        self.sensorSize      = sensorSize
        self.volF = []
        self.volS = []
        self.volP = []
        self.vecT = [] # time vector

        print("SnapShots Field="+str(Field)+" enableNumPy="+str(enableNumPy)+" enableCsv="+str(enableCsv)+" enableVolume="+str(enableVolume)+ " enableSignals="+str(enableSignals)+ " enableMaterial="+str(enableMaterial))
        # if self.File_path does not exist, create it
        if not os.path.exists(self.File_path):
            os.makedirs(self.File_path)
              
    def save_fields(self, Fx, Fy, n):

        xn = Fx.shape[1]
        yn = Fx.shape[0]
        print('['+str(int(n / self.Step))+'] Saving fields x/y '+Label+' shape=('+str(xn)+','+str(yn)+')')
        
        if self.enableVolume:
          self.volF.append(np.array([Fx,Fy]))
        elif self.enableNumPy:
          FILE = self.Filename + ("-fields-%05d" % (int(n / self.Step)))          
          FILE_PATH = os.path.join(self.File_path, FILE)
          np.save(FILE_PATH,np.array([Fx,Fy]))
          print("Fields saved in numpy format at: ", FILE)
        else:  
          dict = {}
          dict['Vx'] = Fx
          dict['Vy'] = Fy
          FILE = self.Filename + ("-fields-%05d" % (int(n / self.Step))) + ".mat"
          savemat(FILE, dict)

    # @blab+
    # save material distribution matrix
    def save_material(self,SimPack,Material):
        if Material == 'I':
          FILE = self.Filename + "-materials-I"
          FILE_PATH = os.path.join(self.File_path, FILE)
          print('Saving scenario material label distribution matrix I ('+FILE+') ['+str(SimPack.Scenario.I.shape[0])+','+str(SimPack.Scenario.I.shape[1])+']')
          np.save(FILE_PATH,SimPack.Scenario.I)
        if Material == 'Iabs':
          FILE = self.Filename + "-materials-Iabs"
          FILE_PATH = os.path.join(self.File_path, FILE)
          print('Saving scenario material label distribution matrix Iabs ('+FILE+') ['+str(SimPack.Scenario.Iabs.shape[0])+','+str(SimPack.Scenario.Iabs.shape[1])+']')
          np.save(FILE_PATH,SimPack.Scenario.Iabs)
        if Material == 'Im':
          FILE = self.Filename + "-materials-Im"
          FILE_PATH = os.path.join(self.File_path, FILE)
          print('Saving simulation material label distribution matrix Im ('+FILE+') ['+str(SimPack.Simulation.Im.shape[0])+','+str(SimPack.Simulation.Im.shape[1])+']')
          np.save(FILE_PATH,SimPack.Simulation.Im)
        if Material == 'RV':
          # Must be recreated from Im, efit2d replaces rho<2 with very large rho value > 1e24
          MRI, NRI = np.shape(SimPack.Simulation.Im)
          Rho = np.zeros((MRI, NRI), dtype=np.float32)
          VL  = np.zeros((MRI, NRI), dtype=np.float32)
          VT  = np.zeros((MRI, NRI), dtype=np.float32)
          for n in range(len(SimPack.Materials)):
            ind = np.nonzero(SimPack.Simulation.Im == SimPack.Materials[n].Label)
            Rho[ind] = SimPack.Materials[n].Rho
            VL[ind] = SimPack.Materials[n].VL
            VT[ind] = SimPack.Materials[n].VT
          # 'air', (1.24, 344, 1e-30), but handled in efit2d as vaccum, i.e. 0
          # 
          ind = np.nonzero(SimPack.Simulation.Im == 255)
          Rho[ind] = 1.24
          VL[ind] = 344
          VT[ind] = 1e-30
          FILE = self.Filename + "-materials-Rho"
          FILE_PATH = os.path.join(self.File_path, FILE)
          print('Saving simulation material distribution  matrix ('+FILE+') ['+str(Rho.shape[0])+','+str(Rho.shape[1])+'] (min='+str(Rho.min())+',max='+str(Rho.max())+')')
          np.save(FILE_PATH,Rho)
          FILE = self.Filename + "-materials-VL"
          FILE_PATH = os.path.join(self.File_path, FILE)
          print('Saving simulation material distribution matrix ('+FILE+') ['+str(VL.shape[0])+','+str(VL.shape[1])+'] (min='+str(VL.min())+',max='+str(VL.max())+')')
          np.save(FILE_PATH,VL)
          FILE = self.Filename + "-materials-VT"
          FILE_PATH = os.path.join(self.File_path, FILE)
          print('Saving simulation material distribution  matrix ('+FILE+') ['+str(VT.shape[0])+','+str(VT.shape[1])+'] (min='+str(VT.min())+',max='+str(VT.max())+')')
          np.save(FILE_PATH,VT)

        
    def save_signal(self, D , n, Label, dt):
    
        self.Label = Label       
        xn = D.shape[1]
        yn = D.shape[0]
        if self.enableVolume:
          print('['+str(int(n / self.Step))+'@'+str(dt*n)+'] Storing signal '+Label+' field shape=('+str(xn)+','+str(yn)+')'+
                ' sensorShape='+str(self.sensorShape))
        else:
          print('['+str(int(n / self.Step))+'@'+str(dt*n)+'] Saving signal '+Label+' field shape=('+str(xn)+','+str(yn)+')'+
                ' sensorShape='+str(self.sensorShape))
        
        dx = 1 # sensor delta
        dy = 1
        ox = 0 # offset
        oy = 0
        if len(self.sensorShape) > 0:
          # size of sensor matrix
          sx = self.sensorShape[0]
          sy = self.sensorShape[1]
          S = np.zeros([sy,sx])
          if len(self.sensorPlacement):
            p=self.sensorPlacement
            ox=p[0]
            dx=p[1]
            oy=p[2]
            dy=p[3]
          else:
            ox=dx=round(xn/(sx+1.0))
            oy=dy=round(yn/(sy+1.0))
          for x in range(0,sx):
            for y in range(0,sy):
              S[y,x]=D[oy+y*dy,ox+x*dx]
        else:
          S = D 
        
        if self.enableVolume:
          self.volS.append(np.copy(S))
          self.vecT.append(dt*n)
        elif self.enableNumPy:
          FILE = self.Filename + ("-signal-"+Label+"-%05d" % (int(n / self.Step)))
          FILE_PATH = os.path.join(self.File_path, FILE)
          np.save(FILE_PATH,S)
        elif self.enableCsv:
          FILE = self.Filename + ("-signal-"+Label+"-%05d" % (int(n / self.Step)))+".csv"
          FILE_PATH = os.path.join(self.File_path, FILE)
          np.savetxt(FILE_PATH, S, delimiter=",")         
        else:
          FILE = self.Filename + ("-signal-"+Label+"-%05d" % (int(n / self.Step))) + ".mat"
          savemat(FILE, S)


    def save_summary (self,dt,dx):
        if self.enableSummary == True:
          FILE = self.Filename + ("-simulation.yml")
          FILE_PATH = os.path.join(self.File_path, FILE)
          # save in yaml format
          data = dict (
            dt=dt,
            dx=dx
          )
          print("Saving simulation summary in "+FILE_PATH)
          with open(FILE_PATH, 'w') as outfile:
            yaml.dump(data, outfile, default_flow_style=False)
        
    def save_power(self, Vx, Vy, n):
       
        S =  np.sqrt(Vx ** 2 + Vy ** 2)

        xn = S.shape[1]
        yn = S.shape[0]
        print('['+str(int(n / self.Step))+'] Saving power '+Label+' shape=('+str(xn)+','+str(yn)+')')
        
        if self.enableVolume:
          self.volP.append(np.copy(S))
        elif self.enableNumPy:
          FILE = self.Filename + ("-power-%05d" % (int(n / self.Step)))
          FILE_PATH = os.path.join(self.File_path, FILE)
          np.save(FILE_PATH,S)
        else:
          FILE = self.Filename + ("-power-%05d" % (int(n / self.Step))) + ".mat"
          FILE_PATH = os.path.join(self.File_path, FILE)
          savemat(FILE_PATH, S)

    def save_fig(self, SV, n, idx=None):

        SV += self.DB
        ind = np.nonzero(SV < 0)
        SV[ind] = 0
        SV /= np.max(SV)

        if self.Color == 0:
            cmap = plt.get_cmap('jet')
        elif self.Color == 1:
            cmap = plt.get_cmap('gray')

        try:
            _resize = False
            M, N = np.shape(SV)
            if M >= 2000:
                _resize = True
            if N >= 2000:
                _resize = True

            if _resize:
                SVV = SV[::3, ::3]
                rgba_img = cmap(SVV)
            else:
                SVV = SV[::2, ::2]
                rgba_img = cmap(SVV)

            rgb_img = np.delete(rgba_img, 3, 2)

            if idx is not None:
                FILE = self.Filename + "_insp%d_" % idx + str(int(n / self.Step)) + self.Extension
            else:
                FILE = self.Filename + ("%05d" % (int(n / self.Step))) + self.Extension
            FILE_PATH = os.path.join(self.File_path, FILE)
            imsave(FILE_PATH, rgb_img)

        except:
          raise ValueError("Too large Snapshots Size!!!!!!!!")

    def save_vol(self):
        
        if len(self.volF) > 0:
          if self.enableNumPy:
            FILE = self.Filename + '-fields'
            FILE_PATH = os.path.join(self.File_path, FILE)
            vol=np.array(self.volF)
            print("Saving numpy S-volume "+str(vol.shape)+" in "+FILE_PATH)
            np.save(FILE_PATH,vol)
            self.volF=[]
            
        if len(self.volS) > 0:
          if self.enableNumPy:
            FILE = self.Filename + '-signal-'+self.Label
            FILE_PATH = os.path.join(self.File_path, FILE)
            vol=np.array(self.volS)
            print("Saving numpy S-volume "+str(vol.shape)+" in "+FILE_PATH)
            np.save(FILE_PATH,vol)
            self.volS=[]
          if self.enableCsv:
            FILE = self.Filename + '-signal-'+self.Label+".csv"
            FILE_PATH = os.path.join(self.File_path, FILE)
            vol=np.array(self.volS)
            vec=np.array(self.vecT)
            vec=vec.reshape(vec.shape[0],1)
            if vol.shape[1]==1:
              vol=vol.reshape(vol.shape[0],vol.shape[1])
            elif vol.shape[2]==1:
              vol=vol.reshape(vol.shape[0],vol.shape[2])
            vol=np.hstack((vec,vol))
            print("Saving reduced CSV S-volume "+str(vol.shape)+" in "+FILE_PATH)
            header="time"
            for i in range(1,vol.shape[1]):
              header=header+",sig"+str(i)
            np.savetxt(FILE_PATH, vol, delimiter=",",header=header,comments="")
            self.volS=[]

        if len(self.volP) > 0:
          if self.enableNumPy:
            FILE = self.Filename + '-power'
            FILE_PATH = os.path.join(self.File_path, FILE)
            vol=np.array(self.volP)
            print("Saving numpy S-volume "+str(vol.shape)+" in "+FILE_PATH)
            np.save(FILE_PATH,vol)
            self.volP=[]
        
