#!/usr/bin/env python

"""
Mapping classes
"""

# ...for the plotting.
import matplotlib.pyplot as plt

# ...for the colour scaling.
from matplotlib import colorbar, colors

# ...for the world map image.
import matplotlib.image as mpimg

from mpl_toolkits.axes_grid1 import make_axes_locatable

# ...for the MATH.
import numpy as np

# ...for the spectra.
from spectra import FluxSpectrum

## ...for the macros.
#from macros import MakeSourceMacro

class MapPoint:

    def __init__(self, lat, lon, alt, day, velx, vely, velz, eSpec, pSpec, dbg = False):

        self.dbg = dbg

        ## Latitude [deg.].
        self.lat = lat

        ## Longitude [deg.].
        self.lon = lon

        ## Altitude [km].
        self.alt = alt

        ## Time [day].
        self.time = day

        ## Velocity x [km s-1]
        self.velx = velx

        ## Velocity y [km s-1]
        self.vely = vely

        ## Velocity z [km s-1]
        self.velz = velz

        ## Electron flux spectrum.
        self.eSpec = eSpec

        ## Proton flux spectrum.
        self.pSpec = pSpec

        if self.dbg:
            print("* DEBUG:")
            print("* DEBUG: Initialising MapPoint object ")
            print("* DEBUG:------------------------------")
            print("* DEBUG: * Latitude : %9.3f [deg.]"   % (self.lat))
            print("* DEBUG: * Longitude: %9.3f [deg.]"   % (self.lon))
            print("* DEBUG: * Altitude : %9.3f [km]"     % (self.alt))
            print("* DEBUG: * Time     : %9.3f [days]"   % (self.time))
            print("* DEBUG:")
            print("* DEBUG: * Velocity : (%f, %f, %f) [km s-1]" % (self.velx, self.vely, self.velz))
            print("* DEBUG:")
            print("* DEBUG: * Total electron flux in [% 8.2f, % 8.2f) MeV: % 10.4f [cm-2 s-1]" % (self.eSpec.GetLowerBound(), self.eSpec.GetUpperBound(), self.eSpec.GetTotalFlux()[0]))
            print("* DEBUG: * Total proton   flux in [% 8.2f, % 8.2f) MeV: % 10.4f [cm-2 s-1]" % (self.pSpec.GetLowerBound(), self.pSpec.GetUpperBound(), self.pSpec.GetTotalFlux()[0]))
            print("* DEBUG:")

    def GetLat(self):
        return self.lat

    def GetLon(self):
        return self.lon

    def GetAlt(self):
        return self.alt

    def GetElectronFlux(self):
        return self.eSpec.GetTotalFlux()[0]

    def GetProtonFlux(self):
        return self.pSpec.GetTotalFlux()[0]

    def MakeElectonSource(self):
        return self.eSpec.MakeSource()

    def MakeProtonSource(self):
        return self.pSpec.MakeSource()


class Map:

    def __init__(self, coordFilePath, attFilePath, eSpecFilePath, pSpecFilePath, e_lo, e_up, p_lo, p_up, outputpath, dbg = False):

        self.dbg = dbg

        self.outputpath = outputpath

        # The minimum altitude for the map.
        self.minalt = 1000000000000. # [km]

        # The maximum altitude for the map.
        self.maxalt = 0.0            # [km]

        ## The colour map for the pixel counts.
        self.cmap = plt.cm.hot # "Hot" used.

        ## The file path of the SPENVIS-generated orbit coordinate file.
        self.coordFilePath = coordFilePath

        ## The file path of the SPENVIS-generated orbit attitude file.
        self.attFilePath = attFilePath

        ## The file path of the SPENVIS-genereted electron flux file.
        self.eSpecFilePath = eSpecFilePath

        ## The file path of the SPENVIS-genereted proton flux file.
        self.pSpecFilePath = pSpecFilePath

        ## The electron spectrum lower bound.
        self.e_lo = e_lo

        ## The electron spectrum upper bound.
        self.e_up = e_up

        ## The proton spectrum lower bound.
        self.p_lo = p_lo

        ## The proton spectrum upper bound.
        self.p_up = p_up

        ##
        self.sphere_radius = 5.0 # [cm].

        ## The geometric factor for the flux calculations.
        self.facA = 0.5 * np.pi * self.sphere_radius * self.sphere_radius # [cm^{2}]


        if self.dbg:
            print("* DEBUG:")
            print("* DEBUG: Initialising Map object                       ")
            print("* DEBUG:-----------------------------------------------")
            print("* DEBUG: * e- spectrum between = [% 8.2f, % 8.2f) MeV" % (self.e_lo, self.e_up))
            print("* DEBUG: * p+ spectrum between = [% 8.2f, % 8.2f) MeV" % (self.p_lo, self.p_up))
            print("* DEBUG: * LUCID sphere radius = % 8.2f [cm]" % (self.sphere_radius))
            print("* DEBUG: * Geometric factor    = % 8.2f [cm^{2}]" % (self.facA))
            print("* DEBUG: *")

        # Read the maps in.
        self.ReadMap()

    def ReadMap(self):
        """
        Reads in the SPENVIS orbit coordinate and attitude files and populates
        the map.
        """

        if self.dbg:
            print("* DEBUG: Calling ReadMap method.")

        ## The coordinate file.
        cf = open(self.coordFilePath, 'r')

        ## The attitude file.
        af = open(self.attFilePath, 'r')

        ## The e flux spectra file.
        ef = open(self.eSpecFilePath, 'r')

        ## The p flux spectra file.
        pf = open(self.pSpecFilePath, 'r')

        ## Read in the files.
        coordlines = cf.readlines()
        attitlines = af.readlines()
        especlines = ef.readlines()
        pspeclines = pf.readlines()

        # Close the input files.
        cf.close()
        af.close()
        ef.close()
        pf.close()

        # Get the spectra details
        #-------------------------

        # Electrons
        e_lEs = especlines[33].split(',')
        #print e_lEs

        # Remove the "Energy" label.
        del e_lEs[0]

        # Get the number of energy bins and remove it.
        e_numEs = int(e_lEs[0])
        del e_lEs[0]

        # Get the energy units.
        e_eUnit = e_lEs[-1]
        del e_lEs[-1]

        self.e_Es = {}

        # Get the actual energy bin edges.
        for i, E in enumerate(e_lEs):
            self.e_Es[i] = float(E.strip())

        if self.dbg:
            print("* DEBUG:")
            print("* DEBUG: Electron energy bins:")
            print("* DEBUG: --> Number of bins = %d (%d)" % (e_numEs, len(e_lEs)))
            #print e_Es

        # Create the spectrum

        # Protons
        p_lEs = pspeclines[33].split(',')
        #print p_lEs

        # Remove the "Energy" label.
        del p_lEs[0]

        # Get the number of energy bins and remove it.
        p_numEs = int(p_lEs[0])
        del p_lEs[0]

        # Get the energy units.
        p_eUnit = p_lEs[-1]
        del p_lEs[-1]

        self.p_Es = {}

        # Get the actual energy bin edges.
        for i, E in enumerate(p_lEs):
            self.p_Es[i] = float(E.strip())

        if self.dbg:
            print("* DEBUG:")
            print("* DEBUG: Proton energy bins:")
            print("* DEBUG: --> Number of bins = %d (%d)" % (p_numEs, len(p_lEs)))
            #print self.p_Es

        self.mappoints = []

        # Process the coordinates.
        for i in range(73, 1514): # hard-coded for now...

            # Extract the coordinates from the line.
            coords = coordlines[i].split(',')

            ## The time [days].
            day = float(coords[0].strip())

            ## The altitude.
            alt = float(coords[1].strip())
            if alt > self.maxalt:
                self.maxalt = alt
            if alt < self.minalt:
                self.minalt = alt

            lat = float(coords[2].strip())

            lon = float(coords[3].strip())
            if lon > 180.0:
                lon = lon - 360.0

            # Information from the attitude file.
            attits = attitlines[i].split(',')

            # Satellite velocities.
            velx = float(attits[ 9].strip())
            vely = float(attits[10].strip())
            velz = float(attits[11].strip())

            # Extract the electron spectra
            #------------------------------
            espec = especlines[i+1].split(',') # +1 because of extra line in file...
            #print espec
            # Remove the B field and L values.
            del espec[0]
            del espec[0]

            ## The integrated flux values.
            e_ifs = {}

            for j, ifs in enumerate(espec):
                e_ifs[j] = float(ifs.strip())

            #print e_ifs
            #if self.dbg:
            #    print("*")
            #    print("* Electron integrated flux values:")
            #    print("*--> Number of values = %d" % (len(e_ifs)))
            #    #print e_ifs

            # Create the electron spectrum
            eSpec = FluxSpectrum("e-", self.e_Es, e_ifs, self.e_lo, self.e_up)

            # Extract the proton spectrum
            #-----------------------------
            pspec = pspeclines[i+1].split(',') # +1 because of extra line in file...
            #print pspec
            # Remove the B field and L values.
            del pspec[0]
            del pspec[0]

            ## The integrated flux values.
            p_ifs = {}

            for j, ifs in enumerate(pspec):
                p_ifs[j] = float(ifs.strip())

            #print p_ifs
            #if self.dbg:
            #    print("*")
            #    print("* Proton integrated flux values:")
            #    print("*--> Number of values = %d" % (len(p_ifs)))
            #    #print p_ifs

            # Create the proton spectrum.
            pSpec = FluxSpectrum("proton", self.p_Es, p_ifs, self.p_lo, self.p_up)

            #mpoint = MapPoint(lat,lon,alt,day,velx,vely,velz,True)
            self.mappoints.append(MapPoint(lat,lon,alt,day,velx,vely,velz,eSpec,pSpec,False))

            #break

    def MakeCoordMap(self):

        self.coordmapfig = plt.figure(101, figsize=(7.0, 4.0), dpi=150, facecolor='w', edgecolor='w')
        #self.coordmapfig.subplots_adjust(bottom=0.08, left=0.08, right=0.8, top=1.0)
        self.cmax = self.coordmapfig.add_subplot(111)
        # y axis
        plt.ylabel('Latitude [\\textdegree]')
        # x axis
        plt.xlabel('Longitude [\\textdegree]')
        #

        # Add the world map background.
        worldimg = mpimg.imread('world_NASA.png')

        #imgplot = plt.imshow(img)
        plt.imshow(worldimg, zorder=0, extent=[-180.0, 180.0, -90., 90.0])

        lats = []
        lons = []
        alts = []

        # Plot the coordinate points - with the altitude.
        for mp in self.mappoints:
            lats.append(mp.GetLat())
            lons.append(mp.GetLon())
            alts.append(mp.GetAlt())

        # Convert the point information to arrays.
        a_lats = np.asarray(lats)
        a_lons = np.asarray(lons)
        a_alts = np.asarray(alts)

        # Create the scatter plot of the flux points.
        ap = plt.scatter(a_lons, a_lats, edgecolors='black', s=10.0, c=a_alts, cmap=self.cmap) #,norm=colors.LogNorm())

        # Divide the axes up to make the colorbar.
        divider = make_axes_locatable(self.cmax)
        self.cax = divider.append_axes("right", size="5%", pad=0.05)

        # Create the colorbar.
        cb = plt.colorbar(ap, cax = self.cax)

        # Add the label.
        cb.set_label('Altitude [km]')

        # Add a grid.
        plt.grid(1)

        # Set the plot limits.
        self.cmax.set_ylim([ -90.0,  90.0])
        self.cmax.set_xlim([-180.0, 180.0])

        # Implement the tight layout.
        plt.tight_layout()

        #plt.show()
        self.coordmapfig.savefig(self.outputpath + "/coords.png")

    def MakeElectronFluxMap(self):

        self.efluxmapfig = plt.figure(102, figsize=(7.0, 4.0), dpi=150, facecolor='w', edgecolor='w')
        #self.efluxmapfig.subplots_adjust(bottom=0.08, left=0.08, right=0.8, top=1.0)
        self.efax = self.efluxmapfig.add_subplot(111)
        # y axis
        plt.ylabel('Latitude [\\textdegree]')
        # x axis
        plt.xlabel('Longitude [\\textdegree]')
        #

        # Add the world map background.
        worldimg = mpimg.imread('world_NASA.png')
        im = self.efax.imshow(worldimg, zorder=0, extent=[-180.0, 180.0, -90., 90.0])

        lats = []
        lons = []
        fluxes = []

        # Plot the coordinate points - with the total electron flux.
        numpoints = 0
        for mp in self.mappoints:
            if mp.GetElectronFlux() > 0.0:
                lats.append(mp.GetLat())
                lons.append(mp.GetLon())
                fluxes.append(mp.GetElectronFlux())
                numpoints += 1

        print("* Number of non-zero electron fluxes = %d" % (numpoints))

        # Convert the point information to arrays.
        a_lats = np.asarray(lats)
        a_lons = np.asarray(lons)
        a_fluxes = np.asarray(fluxes)

        # Create the scatter plot of the flux points.
        fp = plt.scatter(a_lons, a_lats, edgecolors='black', s=10.0, c=a_fluxes, cmap=self.cmap,norm=colors.LogNorm())

        # Divide the axes up to make the colorbar.
        divider = make_axes_locatable(self.efax)
        self.cax = divider.append_axes("right", size="5%", pad=0.05)

        # Create the colorbar.
        cb = plt.colorbar(fp, cax = self.cax)

        # Add the label.
        cb.set_label('$\Phi_{e}$ [$\\textrm{cm}^{-2} \\, \\textrm{s}^{-1}$]')

        # Add a grid.
        plt.grid(1)

        # Set the plot limits.
        self.efax.set_ylim([ -90.0,  90.0])
        self.efax.set_xlim([-180.0, 180.0])

        # Implement the tight layout.
        plt.tight_layout()

        # Show the plot.
        #plt.show()
        self.efluxmapfig.savefig(self.outputpath + "/electronflux.png")

    def MakeProtonFluxMap(self):

        self.pfluxmapfig = plt.figure(103, figsize=(7.0, 4.0), dpi=150, facecolor='w', edgecolor='w')
        #self.pfluxmapfig.subplots_adjust(bottom=0.08, left=0.08, right=0.8, top=1.0)
        self.pfax = self.pfluxmapfig.add_subplot(111)
        # y axis
        plt.ylabel('Latitude [\\textdegree]')
        # x axis
        plt.xlabel('Longitude [\\textdegree]')
        #

        # Add the world map background.
        worldimg = mpimg.imread('world_NASA.png')
        im = self.pfax.imshow(worldimg, zorder=0, extent=[-180.0, 180.0, -90., 90.0])

        lats = []
        lons = []
        fluxes = []

        # Plot the coordinate points - with the total proton flux.
        numpoints = 0
        for mp in self.mappoints:
            if mp.GetProtonFlux() > 0.0:
                lats.append(mp.GetLat())
                lons.append(mp.GetLon())
                fluxes.append(mp.GetProtonFlux())
                numpoints += 1

        print("* Number of non-zero proton fluxes = %d" % (numpoints))

        # Convert the point information to arrays.
        a_lats = np.asarray(lats)
        a_lons = np.asarray(lons)
        a_fluxes = np.asarray(fluxes)

        # Create the scatter plot of the flux points.
        fp = plt.scatter(a_lons, a_lats, edgecolors='black', s=10.0, c=a_fluxes, cmap=self.cmap,norm=colors.LogNorm())

        # Divide the axes up to make the colorbar.
        divider = make_axes_locatable(self.pfax)
        self.cax = divider.append_axes("right", size="5%", pad=0.05)

        # Create the colorbar.
        cb = plt.colorbar(fp, cax = self.cax)

        # Add the label.
        cb.set_label('$\Phi_{p}$ [$\\textrm{cm}^{-2} \\, \\textrm{s}^{-1}$]')

        # Add a grid.
        plt.grid(1)

        # Set the plot limits.
        self.pfax.set_ylim([ -90.0,  90.0])
        self.pfax.set_xlim([-180.0, 180.0])

        # Implement the tight layout.
        plt.tight_layout()

        # Show the plot.
        #plt.show()
        self.pfluxmapfig.savefig(self.outputpath + "/protonflux.png")

    def MakeNumParticlesMap(self, acqtime):

        self.npmapfig = plt.figure(104, figsize=(7.0, 4.0), dpi=150, facecolor='w', edgecolor='w')
        #self.npmapfig.subplots_adjust(bottom=0.08, left=0.08, right=0.8, top=1.0)
        self.npax = self.npmapfig.add_subplot(111)
        # y axis
        plt.ylabel('Latitude [\\textdegree]')
        # x axis
        plt.xlabel('Longitude [\\textdegree]')
        #

        # Add the world map background.
        worldimg = mpimg.imread('world_NASA.png')
        im = self.npax.imshow(worldimg, zorder=0, extent=[-180.0, 180.0, -90., 90.0])

        lats = []
        lons = []

        ## The total number of particles to simulate in each frame.
        nps = []

        # Plot the coordinate points - with the total proton flux.
        numpoints = 0
        for mp in self.mappoints:
            if mp.GetProtonFlux() > 0.0 or mp.GetElectronFlux() > 0.0:
                lats.append(mp.GetLat())
                lons.append(mp.GetLon())
                numberofparticlesp = self.facA * (mp.GetProtonFlux() + mp.GetElectronFlux()) * acqtime
                nps.append(numberofparticlesp)
                numpoints += 1

        print("* Number of populated points = %d" % (numpoints))

        # Convert the point information to arrays.
        a_lats = np.asarray(lats)
        a_lons = np.asarray(lons)
        a_nps  = np.asarray(nps)

        # Create the scatter plot of the map points.
        npp = plt.scatter(a_lons, a_lats, edgecolors='black', s=10.0, c=a_nps, cmap=self.cmap,norm=colors.LogNorm())

        # Divide the axes up to make the colorbar.
        divider = make_axes_locatable(self.npax)
        self.cax = divider.append_axes("right", size="5%", pad=0.05)

        # Create the colorbar.
        cb = plt.colorbar(npp, cax = self.cax)

        # Add the label.
        cb.set_label('Particles per frame')

        # Add a grid.
        plt.grid(1)

        # Set the plot limits.
        self.npax.set_ylim([ -90.0,  90.0])
        self.npax.set_xlim([-180.0, 180.0])

        # Implement the tight layout.
        plt.tight_layout()

        # Show the plot.
        #plt.show()

        self.npmapfig.savefig(self.outputpath + "/numpar.png")

    def MakeSourceMacros(self, acqTime, numPartPerFile):

        # Loop over the map points.
        for mp in self.mappoints:
            if mp.GetProtonFlux() > 0.0 or mp.GetElectronFlux() > 0.0:

                sm = MakeSourceMacro(mp, numPartPerFile)

                #if self.dbg:
                #    print("* DEBUG: source macro")
                #    print sm
                sfn = "lucid-001_lat%+07.2fdeg_lon%+07.2fdeg_alt%07.2fkm_%06.2fs.in" % (mp.GetLat(), mp.GetLon(), mp.GetAlt(), acqTime)
                sf = open(self.outputpath + "/" + sfn, 'w')
                sf.write(sm)
                sf.close()

def MakeSourceMacro(mp, numPartPerFile):
    s = ""
    s += "#=============================================================================\n"
    s += "#       Allpix: simulation of the LUCID experiment (with visualisation)       \n"
    s += "#=============================================================================\n"
    s += "\n"
    s += "# The laboratory\n"
    s += "# Outer space!\n"
    s += "/allpix/lab/setLatitude   %10.6f deg\n" % (mp.GetLat())
    s += "/allpix/lab/setLongitude  %10.6f deg\n" % (mp.GetLon())
    s += "/allpix/lab/setAltitude   %10.5f km\n"  % (mp.GetAlt())
    s += "/allpix/lab/setRollAngle  0.0 deg\n"
    s += "/allpix/lab/setPitchAngle 0.0 deg\n"
    s += "/allpix/lab/setYawAngle   0.0 deg\n"
    s += "\n"
    s += "# The LUCID TPX0 detector \n"
    s += "/allpix/det/setId        10000\n"
    s += "/allpix/det/setPosition   0.0   18.3  18.3 mm\n"
    s += "/allpix/det/setRotation   0.0    0.0   0.0 deg\n"
    s += "/allpix/det/setFoilThickness 0.0174 mm\n"
    s += "\n"
    s += "# The LUCID TPX1 detector\n"
    s += "/allpix/det/setId        10001\n"
    s += "/allpix/det/setPosition  18.3   18.3   0.0 mm\n"
    s += "/allpix/det/setRotation   0.0  -90.0   0.0 deg\n"
    s += "/allpix/det/setFoilThickness 0.0174 mm\n"
    s += "\n"
    s += "# The LUCID TPX2 detector.\n"
    s += "/allpix/det/setId        10002\n"
    s += "/allpix/det/setPosition -18.3   18.3   0.0 mm\n"
    s += "/allpix/det/setRotation   0.0   90.0   0.0 deg\n"
    s += "/allpix/det/setFoilThickness 0.0174 mm\n"
    s += "\n"
    s += "# The LUCID TPX3 detector \n"
    s += "/allpix/det/setId        10003\n"
    s += "/allpix/det/setPosition   0.0   18.3 -18.3 mm\n"
    s += "/allpix/det/setRotation   0.0  180.0   0.0 deg\n"
    s += "/allpix/det/setFoilThickness 0.0174 mm\n"
    s += "\n"
    s += "# The LUCID TPX4 detector \n"
    s += "/allpix/det/setId        10004\n"
    s += "/allpix/det/setPosition   -3.01   -7.600   0.0 mm\n"
    s += "/allpix/det/setRotation   90.0     0.0    90.0   deg\n"
    s += "/allpix/det/setFoilThickness 0.0174 mm\n"
    s += "\n"
    s += "# The test structure - the LUCID experiment.\n"
    s += "/allpix/extras/setTestStructure LUCID\n"
    s += "/allpix/extras/setTestStructurePosition 0.0 0.0 0.0 mm\n"
    s += "/allpix/extras/setTestStructureRotation 0.0 0.0 0.0 deg\n"
    s += "/allpix/extras/setTestStructureShieldingThickness 0.68 mm\n"
    s += "\n"
    s += "# We're in space, of course.\n"
    s += "/allpix/extras/setWorldMaterial Vacuum\n"
    s += "\n"
    s += "# The maximum step length (applies only to the sensor).\n"
    s += "/allpix/det/setMaxStepLengthSensor 10 um\n"
    s += "\n"
    s += "# Initialize the simulation.\n"
    s += "/run/initialize\n"
    s += "\n"
    s += "# Build the detectors.\n"
    s += "/allpix/det/update\n"
    s += "\n"
    s += "\n"
    s += "##############################################################################\n"
    s += "\n"
    s += "/tracking/verbose 0\n"
    s += "\n"
    s += "##############################################################################\n"
    s += "#                       General Particle Source (GPS)\n"
    s += "##############################################################################\n"
    s += "\n"

    # Get the electron energy histogram (if required).
    if mp.GetElectronFlux() > 0.0:
        s += mp.MakeElectonSource()

    if mp.GetProtonFlux() > 0.0 and mp.GetProtonFlux() > 0.0:
        s += "#\n"
        s += "/gps/source/add % 15.5f\n" % (mp.GetProtonFlux())
        s += "#\n"

    # Get the proton energy histogram (if required).
    if mp.GetProtonFlux() > 0.0:
        s += mp.MakeProtonSource()

    s += "#\n"
    s += "/allpix/lab/setSourceId LUCID, e- and p+\n"
    s += "\n"
    s += "\n"
    s += "/gps/source/list\n"
    s += "\n"
    s += "##############################################################################\n"
    s += "\n"
    s += "# Run the simulation with the number of events specified.\n"
    s += "/run/beamOn %d\n" % (numPartPerFile)
    s += "\n"
    s += "##############################################################################\n"

    return s