#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Oct 16 09:36:43 2023

@author: Phillip Manley
"""

import os
import numpy as np
from scipy import constants
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.ticker as ticker
import jcmwave as jcm
from matplotlib.patches import Rectangle
from IPython.display import HTML, clear_output
try:
    from dipole import electric_dipole_efield_full, electric_dipole_bfield_full
except ImportError:
    from .dipole import electric_dipole_efield_full, electric_dipole_bfield_full

comp_names = {0:'x', 1:'y', 2:'z'}

class DipoleNearFieldVisualization(object):

    
    def __init__(self, single_frame=False, case='analytical'):
        self.setup_fig_axes()
        self.single_frame = single_frame
        
        #animation settings
        self.fps = 30
        self.interval = 100 #ms
        self.total_time = 3. #seconds
        self.n_frames = int(self.fps*self.total_time)        
        self.case = case
        self.setup_fig_axes()
        self.simulation_keys = self.init_simulation_keys()
        self.e_field = None
        self.b_field = None
        self.setup_subfigures()
        #keys = self.simulation_keys()
        
        
        

    def setup_fig_axes(self):
        plt.close("all")
        fig, ax = plt.subplots(2, 3, figsize=(20, 12), gridspec_kw={'wspace':0.5, 'hspace':0.5})    
        self.fig = fig
        self.axes = ax            
        
    def init_simulation_keys(self):
        omega = 1e15
        
        self.time_array = np.linspace(0, np.pi*2., self.n_frames, endpoint=False)/omega        
        
        p = np.array([1., 0., 0])*constants.e/omega
        keys = {}
        keys['omega'] = omega
        keys['p'] = p
        keys['dipole_pos'] = np.array([0., 0., 0.])
        keys['efield_comp'] = 0
        keys['hfield_comp'] = 1
        keys['plane_1'] = {}
        n_points = 101
        keys['plane_1']['x'] = np.linspace(-1, 1, n_points)*2000e-9
        keys['plane_1']['y'] = np.array([1e-12])
        keys['plane_1']['z'] = np.linspace(-1, 1, n_points)*2000e-9
        
        keys['plane_2'] = {}
        keys['plane_2']['x'] = np.array([1e-12])
        keys['plane_2']['y'] = np.linspace(-1, 1, n_points)*2000e-9
        keys['plane_2']['z'] = np.linspace(-1, 1, n_points)*2000e-9
        
        keys['plane_3'] = {}
        keys['plane_3']['x'] = np.linspace(-1, 1, n_points)*2000e-9
        keys['plane_3']['y'] = np.linspace(-1, 1, n_points)*2000e-9
        keys['plane_3']['z'] = np.array([1e-12])
        
        keys['t'] = 0.
        return keys
           
    def get_e_field(self, plane, r, array_shape):
        if self.case == 'analytical':
            return self.get_e_field_analytical(r, array_shape)
        elif self.case == 'numerical':
            return self.get_e_field_numerical(plane, "dipole_project")
        elif self.case == 'numerical_scatterer':
            return self.get_e_field_numerical(plane, "dipole_with_scatterer")
        else:
            raise ValueError("unknown case: {}".format(self.case))
    
    def get_e_field_analytical(self, r, array_shape):
        keys = self.simulation_keys
        e_field = electric_dipole_efield_full(keys['omega'], keys['p'],
                                              keys['dipole_pos'], r, time=keys['t'])    
        e_field = e_field.reshape(array_shape + [3])
        return e_field
    
    def get_e_field_numerical(self, plane, project):
        folder = os.path.join("..", "sources", project, "project_results")        
        if plane == 0:
            e_field = jcm.loadcartesianfields(os.path.join(folder ,"e_field_xz.jcm"))['field'][0]
        elif plane == 1:
            e_field = jcm.loadcartesianfields(os.path.join(folder ,"e_field_yz.jcm"))['field'][0]
        elif plane == 2:
            e_field = jcm.loadcartesianfields(os.path.join(folder ,"e_field_xy.jcm"))['field'][0]
        return e_field
    
    def process_e_field(self, e_field):
        keys = self.simulation_keys
        e_field_phase = np.exp(-1j*keys['omega']*keys['t'])*e_field
        e_field_phase = np.real(e_field_phase)        
        e_field_phase = e_field_phase[..., keys['efield_comp']]        
        return e_field_phase
    
    def get_b_field(self, plane, r, array_shape):
        if self.case == 'analytical':
            return self.get_b_field_analytical(r, array_shape)
        elif self.case == 'numerical':
            return self.get_b_field_numerical(plane, "dipole_project")
        elif self.case == 'numerical_scatterer':
            return self.get_b_field_numerical(plane, "dipole_with_scatterer")
        else:
            raise ValueError("unknown case: {}".format(self.case))
    
    def get_b_field_analytical(self, r, array_shape):
        keys = self.simulation_keys
        b_field = electric_dipole_bfield_full(keys['omega'], keys['p'],
                                              keys['dipole_pos'], r, time=keys['t'])    
        b_field = b_field.reshape(array_shape + [3])                
        return b_field
    
    def get_b_field_numerical(self, plane, project):
        folder = os.path.join("..", "sources", project, "project_results")        
        if plane == 0:
            h_field = jcm.loadcartesianfields(os.path.join(folder ,"h_field_xz.jcm"))['field'][0]
        elif plane == 1:
            h_field = jcm.loadcartesianfields(os.path.join(folder ,"h_field_yz.jcm"))['field'][0]
        elif plane == 2:
            h_field = jcm.loadcartesianfields(os.path.join(folder ,"h_field_xy.jcm"))['field'][0]
        b_field = h_field*constants.mu_0
        return b_field
    
    def process_b_field(self, b_field):
        keys = self.simulation_keys
        b_field_phase = np.exp(-1j*keys['omega']*keys['t'])*b_field
        b_field_phase = np.real(b_field_phase)
        b_field_phase = b_field_phase[..., keys['hfield_comp']]
        return b_field_phase
    
    def annotate(self, plane, ax):
        if plane == 0 and self.case == "numerical_scatterer":
            width = 400.
            height = 1000.
            rect = Rectangle((-750.-width*0.5, 750.-height*0.5), width, height, angle=45., rotation_point='center',
                             fill=False, edgecolor = 'k', lw=2.0)
            ax.add_artist(rect)
    
    def setup_subfigures(self):
        keys = self.simulation_keys        
        #print(self.axes.shape)
        self.imgs = np.empty(self.axes.shape, dtype=object) 
        self.e_fields = []
        self.b_fields = []
        self.colorbars = np.empty(self.axes.shape, dtype=object)
        for plane in range(3):
            x = keys['plane_{}'.format(plane+1)]['x']
            y = keys['plane_{}'.format(plane+1)]['y']
            z = keys['plane_{}'.format(plane+1)]['z']
            X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
            
            if plane == 0:
                plot_x = np.squeeze(X)*1e9
                plot_y = np.squeeze(Z)*1e9
            elif plane == 1:
                plot_x = np.squeeze(Y)*1e9
                plot_y = np.squeeze(Z)*1e9
            elif plane == 2:
                plot_x = np.squeeze(X)*1e9
                plot_y = np.squeeze(Y)*1e9            
            
            array_shape = list(np.squeeze(X).shape)
            r = np.vstack([X.flatten(), Y.flatten(), Z.flatten()]).T            
            
            e_field = self.get_e_field(plane, r, array_shape)
            self.e_fields.append(e_field)
            e_field = self.process_e_field(e_field)
            
            #plt.sca(self.axes[0, plane])
            ax = self.axes[0, plane]
            vmax = 2e-5
            vmin = -vmax            
            
            img = ax.pcolormesh(plot_x, plot_y, e_field, vmin=vmin, vmax=vmax, cmap='coolwarm',
                                shading='gouraud')
            self.imgs[0, plane] = img
            cb1 = self.fig.colorbar(img, ax=ax)
            self.colorbars[0, plane] = cb1
            self.annotate(plane, ax)
            
            
            b_field = self.get_b_field(plane, r, array_shape)                            
            self.b_fields.append(b_field)
            b_field = self.process_b_field(b_field)
            h_field = b_field/constants.mu_0
            
            #plt.sca(self.axes[1, plane])
            ax = self.axes[1, plane]
            vmax = 1e-7
            vmin = -vmax
            
            img = ax.pcolormesh(plot_x, plot_y, h_field, vmin=vmin, vmax=vmax, cmap='PiYG',
                                shading='gouraud')
            
            self.imgs[1, plane] = img                        
            #cb2 = plt.colorbar()            
            cb2 = self.fig.colorbar(img, ax=ax)
            self.colorbars[1, plane] = cb2
            self.annotate(plane, ax)
        
    
    def init_animation(self):
        keys = self.simulation_keys
        for plane in range(3):
            if plane == 0:            
                x_name = 'X'
                y_name = 'Z'
            elif plane == 1:                
                x_name = 'Y'
                y_name = 'Z'
            elif plane == 2:                
                x_name = 'X'
                y_name = 'Y'            
            
            #plt.sca(self.axes[0, plane])            
            ax = self.axes[0, plane]
            
            ax.set_xlabel("{} (nm)".format(x_name))
            ax.set_ylabel("{} (nm)".format(y_name))
            
            unit_str = "$\Re${E$_{\mathrm{" + "{}".format(comp_names[keys['efield_comp']]) + "}}$} (V/m)"        
            cb1 = self.colorbars[0, plane]
            cb1.ax.set_title(unit_str)            
            formatter = ticker.ScalarFormatter()
            formatter.set_scientific(True)
            formatter.set_powerlimits((-3,3))
            cb1.ax.yaxis.set_major_formatter(formatter)            
            
            #plt.sca(self.axes[1, plane])            
            ax = self.axes[1, plane]
            ax.set_xlabel("{} (nm)".format(x_name))
            ax.set_ylabel("{} (nm)".format(y_name))

            cb2 = self.colorbars[1, plane]
            unit_str = "$\Re${H$_{\mathrm{" + "{}".format(comp_names[keys['hfield_comp']]) + "}}$} (A/m)"
            cb2.ax.set_title(unit_str)    
        return self.imgs.flatten().tolist()

    def update_animation(self, frame):
        keys = self.simulation_keys
        keys['t'] = self.time_array[frame]
        #print("updating with frame: {}".format(frame))
        for plane in range(3):                  
            e_field = self.e_fields[plane]
            e_field = self.process_e_field(e_field)
            self.imgs[0, plane].set_array(e_field)                        
            
            b_field = self.b_fields[plane]
            b_field = self.process_b_field(b_field)        
            h_field = b_field/constants.mu_0            
            
            self.imgs[1, plane].set_array(h_field)
        
        return self.imgs.flatten().tolist()
    
    def animation(self):    
        return animation.FuncAnimation(self.fig, self.update_animation, frames=self.n_frames,
                                           interval=self.interval,
                                           init_func=self.init_animation,
                                           repeat=False,
                                           blit=True)
    
    def to_file(self, file_name="near_fields"):
        anim = self.animation()
        file_name += "_{}".format(self.case)
        file_name += ".mp4"
        full_file_path = os.path.join("..", "videos", file_name)
        anim.save(full_file_path, writer="ffmpeg", fps=self.fps)
        
    def show(self):            
        if self.single_frame:
            self.init_animation()
            return self.fig
        else:
            anim = self.animation()            
            clear_output()
            return HTML(anim.to_jshtml())
        
    
    
if __name__ == "__main__":
    single_frame = False
    vis = DipoleNearFieldVisualization(single_frame=single_frame, case='analytical')
    if single_frame:
        anim = vis.show()    
    else:
        vis.to_file()
    #
    #anim.show()
        
        
    
