#!/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.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.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 DipoleFarFieldVisualizationPolar(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
        style_keys['legend.fontsize'] = 16
        
        mpl.style.use(style_keys)                
        plt.close("all")
        fig, ax = plt.subplots(1, 3, figsize=(16, 6), gridspec_kw={'wspace':0.5, 'hspace':0.5},
                               subplot_kw={'projection': 'polar'})
        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_1'] = {}
        keys['plane_1']['r'] = np.array([1.])
        keys['plane_1']['theta'] = np.radians(np.arange(-180, 181, 1))
        keys['plane_1']['phi'] =np.array([0.])
        
        keys['plane_2'] = {}
        keys['plane_2']['r'] = np.array([1.])
        keys['plane_2']['theta'] = np.radians(np.arange(-180, 181, 1))
        keys['plane_2']['phi'] = np.radians(np.array([90.]))
        
        keys['plane_3'] = {}
        keys['plane_3']['r'] = np.array([1.])
        keys['plane_3']['theta'] = np.radians(np.array([90.]))
        keys['plane_3']['phi'] = np.radians(np.arange(0, 361, 1))
        
        for plane in range(3):
            plane_str = "plane_{}".format(plane+1)
            
            R, THETA, PHI = np.meshgrid(keys[plane_str]['r'],
                                        keys[plane_str]['theta'],
                                        keys[plane_str]['phi'],
                                        indexing='ij')
            X, Y, Z = spherical2cart(R, THETA, PHI)
            X = X.reshape(R.shape)
            Y = Y.reshape(R.shape)
            Z = Z.reshape(R.shape)            
            keys[plane_str]['THETA'] = THETA
            keys[plane_str]['PHI'] = PHI
            array_shape = list(np.squeeze(X).shape)
            keys[plane_str]['array_shape'] = array_shape
            positions = np.vstack([X.flatten(), Y.flatten(), Z.flatten()]).T
            keys[plane_str]['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):
        keys = self.simulation_keys
        plane = keys['plane']
        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 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")        
        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 setup_subfigures(self):
        keys = self.simulation_keys        
        
        for plane in range(3):
            plane_str = "plane_{}".format(plane+1)
            ax = self.axes[plane]
            keys['positions'] = keys[plane_str]['positions']
            keys['array_shape'] = keys[plane_str]['array_shape']
            
            keys['plane'] = plane            
            
            if plane == 0:
                polar_angle = np.squeeze(keys[plane_str]['THETA'])
                title = 'XZ Plane'
            if plane == 1:
                polar_angle = np.squeeze(keys[plane_str]['THETA'])
                title = 'YZ Plane'
            if plane == 2:
                polar_angle = np.squeeze(keys[plane_str]['PHI'])
                title = 'XY Plane'
            ax.set_title(title)
            
            if plane < 2:
                ax.set_theta_offset(np.pi*0.5)
                ax.set_xticks(np.linspace(-np.pi, np.pi, 9)[1:])
                ax.set_thetalim(np.pi, -np.pi)
                ax.set_theta_direction(-1)
                angle = np.pi*0.5*0.2
                ax.text(angle, 0.12*1.15, r"$\rightarrow \theta$", fontsize=18., rotation=np.degrees(-angle), ha='center', va='center')
            else:
                angle = np.pi*0.5*0.2
                ax.text(angle, 0.12*1.15, r"$\leftarrow \phi$", fontsize=18., rotation=np.degrees(angle-np.pi*0.5), ha='center', va='center')
            
            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']            
                
                
                if i_polarization == 0:                
                    ls = '-'
                if i_polarization == 1:                
                    ls = '--'
                if i_polarization == 2:                
                    ls = '-.'
                
                label = "{} Pol.".format(pol_names[i_polarization].upper())
    
                            
    
    
                e_field = self.get_e_field()                
                #e_field = self.process_e_field(e_field)                
                b_field = self.get_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.set_rticks([])
                ax.plot(polar_angle, p_r_real, label=label, ls=ls)
                
                if i_polarization == 2:
                    if plane < 2:
                        ax.legend(bbox_to_anchor=[1.0, 1.0])
                    else:
                        ax.legend(bbox_to_anchor=[1.1, 1.0])
            
        
    def show(self):            
        self.fig.show()        
    
    
if __name__ == "__main__":    
    vis = DipoleFarFieldVisualizationPolar(case='numerical')    
    vis.show()    
    
        
        
    
