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

@author: Phillip Manley
"""

import sys
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
from matplotlib.patches import Rectangle
from IPython.display import HTML, clear_output
import matplotlib as mpl
from matplotlib import colormaps
from matplotlib import colors
import matplotlib.animation as animation
import matplotlib.ticker as ticker

jcm_root = os.getenv('JCMROOT')
if jcm_root is None:
    raise EnvironmentError('No environment variable JCMROOT found. '+
                           'Maybe JCMsuite is not installed.')
sys.path.insert(0,os.path.join(jcm_root, 'ThirdPartySupport', 'Python'))                      
import jcmwave as jcm

try:
    from utils import spherical2cart, cart2spherical
    from dipole import electric_dipole_efield_far, electric_dipole_bfield_far, electric_dipole_strength, electric_dipole_emission
except ImportError:
    from .utils import spherical2cart, cart2spherical
    from .dipole import electric_dipole_efield_far, electric_dipole_bfield_far, electric_dipole_strength, electric_dipole_emission

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


class DipoleFarFieldVisualization3D(object):

    
    def __init__(self, case='analytical'):
        self.setup_fig_axes()
        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()
        

    def setup_fig_axes(self):
        style_keys = {}
        style_keys['axes.titlesize'] = 18
        style_keys['axes.labelsize'] = 18
        style_keys['lines.linewidth'] = 2
        style_keys['xtick.labelsize'] = 16
        style_keys['ytick.labelsize'] = 16
        
        mpl.style.use(style_keys)                
        plt.close("all")
        fig, ax = plt.subplots(1, 3, figsize=(16, 6), gridspec_kw={'wspace':0.3},
                       subplot_kw={'projection': '3d'})
        self.fig = fig
        self.axes = ax            
        
    def init_simulation_keys(self):
        omega = 1e15                
        keys = {}
        keys['omega'] = omega
        power_emission = 1
        keys['dipole_strength'] = electric_dipole_strength(omega, power_emission)        
        
        keys['dipole_pos'] = np.array([0., 0., 0.])        
        
        keys['plane'] = {}        
        delta_theta = 1
        delta_phi = 5
        #dA = np.sin()

        r = np.array([1.]) 
        theta = np.radians(np.arange(0, 180, delta_theta))
        phi = np.radians(np.arange(0, 365, delta_phi))
        
        R, THETA, PHI = np.meshgrid(r, theta, phi, indexing='ij')
        keys['THETA'] = THETA
        keys['PHI'] = PHI
        
        X, Y, Z = spherical2cart(R, THETA, PHI)        
        X = X.reshape(R.shape)
        Y = Y.reshape(R.shape)
        Z = Z.reshape(R.shape)
        
        array_shape = list(np.squeeze(X).shape)
        positions = np.vstack([X.flatten(), Y.flatten(), Z.flatten()]).T
        keys['array_shape'] = array_shape
        keys['positions'] = positions
        return keys
           
    def get_e_field(self):
        if self.case == 'analytical':
            return self.get_e_field_analytical()
        elif self.case == 'numerical':
            return self.get_e_field_numerical("dipole_project")
        elif self.case == 'numerical_scatterer':
            return self.get_e_field_numerical("dipole_with_scatterer")
        else:
            raise ValueError("unknown case: {}".format(self.case))
    
    def get_e_field_analytical(self):
        keys = self.simulation_keys
        e_field = electric_dipole_efield_far(keys['omega'], keys['dipole_pol'],
                                              keys['dipole_pos'], keys['positions'])
        e_field = e_field.reshape(keys['array_shape'] + [3])
        return e_field
    
    def get_e_field_numerical(self, project):        
        folder = os.path.join("..", "sources", project, "project_results")
        file = os.path.join(folder ,"e_field.jcm")
        if not os.path.isfile(file):
            raise IOError("No results available")
        e_field = jcm.loadcartesianfields(file)['field'][0]              
        return e_field
    
    def process_e_field(self, e_field):
        keys = self.simulation_keys        
        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):
        if self.case == 'analytical':
            return self.get_b_field_analytical()
        elif self.case == 'numerical':
            return self.get_b_field_numerical("dipole_project")
        elif self.case == 'numerical_scatterer':
            return self.get_b_field_numerical("dipole_with_scatterer")
        else:
            raise ValueError("unknown case: {}".format(self.case))
    
    def get_b_field_analytical(self):
        keys = self.simulation_keys
        b_field = electric_dipole_bfield_far(keys['omega'], keys['dipole_pol'],
                                              keys['dipole_pos'], keys['positions'])
        b_field = b_field.reshape(keys['array_shape'] + [3])                
        return b_field
    
    def get_b_field_numerical(self, project):
        keys = self.simulation_keys
        plane = keys['plane']
        folder = os.path.join("..", "sources", project, "project_results")
        file = os.path.join(folder ,"h_field.jcm")
        if not os.path.isfile(file):
            raise IOError("No results available")
        h_field = jcm.loadcartesianfields(file)['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 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 i_polarization in range(3):
            
            if i_polarization == 0:
                keys['dipole_pol'] = np.array([1., 0., 0.])*keys['dipole_strength']
            elif i_polarization == 1:
                keys['dipole_pol'] = np.array([0., 1., 0.])*keys['dipole_strength']
            elif i_polarization == 2:
                keys['dipole_pol'] = np.array([0., 0., 1.])*keys['dipole_strength']            
            
            
            x_name = 'X'
            y_name = 'Y'
            z_name = 'Z'
            
            
            e_field = self.get_e_field()
            self.e_fields.append(e_field)
            #e_field = self.process_e_field(e_field)            

            b_field = self.get_b_field()                            
            self.b_fields.append(b_field)
            #b_field = self.process_b_field(b_field)
            h_field = b_field/constants.mu_0

            poynting = 0.5*np.cross(e_field, h_field)
            
            p_r_real, p_theta_real, p_phi_real = cart2spherical(np.real(poynting[..., 0]),
                                                                np.real(poynting[..., 1]),
                                                                np.real(poynting[..., 2]))


            ax = self.axes[i_polarization]            
            vmax = np.amax(p_r_real)*1.8
            vmin = np.amin(p_r_real)
            
            
            ax.set_xlabel(x_name, labelpad=-10.)
            ax.set_ylabel(y_name, labelpad=-10.)
            ax.set_zlabel(z_name, labelpad=-10.)
            ax.yaxis.set_major_formatter(ticker.NullFormatter())
            ax.xaxis.set_major_formatter(ticker.NullFormatter())
            ax.zaxis.set_major_formatter(ticker.NullFormatter())
            
            p_r_x, p_r_y, p_r_z = spherical2cart(p_r_real, keys['THETA'], keys['PHI'])
            plot_x = np.squeeze(p_r_x)
            plot_y = np.squeeze(p_r_y)
            plot_z = np.squeeze(p_r_z)
            ax.set_title("Dipole\nPolarization {}".format(pol_names[i_polarization].upper()))
            
            cmap = colormaps.get_cmap("viridis_r")        
            norm = colors.Normalize(vmin=vmin, vmax=vmax)
            facecolors = cmap(norm(p_r_real))
            
            img = ax.plot_surface(plot_x, plot_y, plot_z, facecolors=facecolors)
            
            ax.axis("equal")
            
            self.imgs[i_polarization] = img
            cb1 = self.fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
                                    ax=ax, shrink=0.5, pad=0.15)
            self.colorbars[i_polarization] = cb1
            
            unit_str = "$\mathbf{S}\cdot\mathbf{n}$ (W)"        
            cb1.ax.set_title(unit_str, pad=20.)
            formatter = ticker.ScalarFormatter()
            formatter.set_scientific(True)
            cb1.ax.yaxis.set_major_formatter(formatter)                   
            
    
    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):            
        self.fig.show()        
    
    
if __name__ == "__main__":    
    vis = DipoleFarFieldVisualization3D(case='numerical')    
    vis.show()    
    
        
        
    
