#!/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 DipoleFarFieldVisualizationPolarComparison(object):

    
    def __init__(self):
        self.case = None
        self.setup_fig_axes()
        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, 1, figsize=(8, 6), subplot_kw={'projection': 'polar'})
        self.fig = fig
        self.axes = ax            
        
    def init_simulation_keys(self):
        omega = 1e15
        keys = {}
        keys['omega'] = omega
        keys['dipole_strength'] = constants.e/omega
        keys['dipole_pol'] = np.array([1., 0., 0.])*keys['dipole_strength']
        keys['dipole_pos'] = np.array([0., 0., 0.])
        
        keys['plane'] = {}
        keys['plane']['r'] = np.array([1.])
        keys['plane']['theta'] = np.radians(np.arange(-180, 181, 2))
        keys['plane']['phi'] = np.array([0.])

        return keys
           
    def get_e_field(self):        
        if self.case == 'analytical':
            return self.get_e_field_analytical()
        elif self.case == 'numerical':
            return self.get_far_field_numerical("dipole_project")
        elif self.case == 'numerical_scatterer':
            return self.get_far_field_numerical("dipole_with_scatterer")
        else:
            raise ValueError("unknown case: {}".format(self.case))
            
    def get_ff(self):
        e_field = self.get_e_field()
        e_norm = np.linalg.norm(np.abs(e_field), axis=1)        
        p = 0.5*e_norm**2/(constants.c*constants.mu_0)
        return p
        
    
    def get_e_field_analytical(self):
        keys = self.simulation_keys
        
        R, THETA, PHI = np.meshgrid(keys['plane']['r'],
                                    keys['plane']['theta'],
                                    keys['plane']['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']['THETA'] = THETA
        keys['plane']['PHI'] = PHI
        array_shape = list(np.squeeze(X).shape)
        keys['plane']['array_shape'] = array_shape
        positions = np.vstack([X.flatten(), Y.flatten(), Z.flatten()]).T
        
        e_field = electric_dipole_efield_far(keys['omega'], keys['dipole_pol'],
                                              keys['dipole_pos'], positions)
        e_field = e_field.reshape(array_shape + [3])
        return e_field
    
    
    def get_far_field_numerical(self, project):        
        folder = os.path.join("sources", project, "project_results")        
        ff = jcm.loadtable(os.path.join(folder ,"ff.jcm"))['ElectricFieldStrength'][0]
        return ff
        
    def setup_subfigures(self):
        keys = self.simulation_keys
        
        color_map = {'analytical': 'tab:blue',
                      'numerical': 'tab:orange',
                      'numerical_scatterer': 'tab:green'}
        
        label_map = {'analytical': 'Analytical',
                     'numerical': 'Numerical',
                     'numerical_scatterer': 'Num. Structured'}
        
        for case in ['analytical', 'numerical', 'numerical_scatterer']:            
            self.case = case
            ax = self.axes
            theta = keys['plane']['theta']            
            ff = self.get_ff()
            if case == 'numerical':
                ax.scatter(theta[::-3], ff[::-3], marker='o', label=label_map[case], c=color_map[case], zorder=2)
            else:
                ax.plot(theta, ff, label=label_map[case], c=color_map[case], lw=5.)
            
        ax.legend(bbox_to_anchor=[1.5, 1.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)
        max_val = 5.194763467937593e-25
        angle = np.pi*0.5*0.2
        ax.text(angle, max_val*1.15, r"$\rightarrow \theta$", fontsize=18., rotation=np.degrees(-angle), ha='center', va='center')
        ax.set_rticks([])
            
    
        
    def show(self):
        return self.fig.show()
    
    
if __name__ == "__main__":    
    vis = DipoleFarFieldVisualizationPolarComparison()    
    vis.show()    
    
        
        
    
