# -*- coding: utf-8 -*-
"""
Created on Sat Apr 29 17:29:05 2017

@author: Ian
"""

import math
import numpy
import Planck
import LineProf
#import matplotlib.pyplot as plt
import Useful


#Helper function - Chapman's integration factor function, Phi(r, mu), Eq. (25)

def Phi(rIndx, mu, xi, t0Indx, it, tau):
    
    #Assume 32-bit to be on the safe side
    tiny = 1.175494e-38 
    tinySq = tiny**2
    huge = 3.4028237e38
    #logTiny = math.log(tiny)
    #logHuge = math.log(huge)
    
    r = xi[rIndx]
    oneMinusMuSq = numpy.double(1.0) - math.pow(mu, 2)
    if (oneMinusMuSq < tiny):
        oneMinusMuSq = tiny
    
    #print("In Phi():")
    #print("r ", r, " mu ", mu, " oneMinusMuSq ", oneMinusMuSq)
    
    PhiRmu = 0.0
    #numer = 0.0
    denom = 0.0
    logNumer = 0.0
    #logDenom = 0.0
    integrand = 0.0
    logIntegrand = 0.0
    integral = 0.0
    
    #integrate *outward*: height xi[i] increases with decreasing i
    for i in range(rIndx, t0Indx, -1):

        deltaXi = xi[i-1] - xi[i]  # should be +ve
        #if (i > 44):
        #    print("i ", i, " logKappa[i] ", logKappa[i], " rho[0][i] ", rho[0][i], " xi[i] ", xi[i])
        #numer = math.exp(logKappa[i]) * rho[0][i]
        #logNumer = logKappa[i] + rho[1][i] + math.log(deltaXi)
        logNumer = math.log(tau[0][i] - tau[0][i-1])
        #numer = math.exp(logNumer)
        #denom = 1.0 - ( ((r/xi[i])**2) * oneMinusMuSq )
        logHelp = numpy.double(2.0)*(math.log(r)-math.log(xi[i])) + math.log(oneMinusMuSq)
        helper = math.exp(logHelp)
        denom = numpy.double(1.0) - helper
        if (denom < tinySq):
            denom = tiny
        else:
            denom = math.sqrt(denom)

        #integrand = numer / denom        
        logIntegrand = logNumer - math.log(denom)
        #logIntegrand = logNumer #Test
        integrand = math.exp(logIntegrand)
        
        #if (i > 44):
            #print("deltaXi ", deltaXi, " numer ", numer, " denom ", denom,\
            #  " integrand ", integrand)
        integral = integral + integrand
        #if (it == 0 and i%10 == 0):
        #    print("i ", i, " logKappa ", logKappa[i], " rho[1] ", rho[1][i],\
        #         " deltaXi ", deltaXi, " integrand*deltaXi ", integrand*deltaXi)
        #if (i > 44):
            #print("integral ", integral)
     
    #if (it == 0):
    #    print("integral ", integral) 
    PhiRmu = math.exp(-1.0 * integral)
    #if (i > 44):
        #print("PhiRmu ", PhiRmu)
    
    return PhiRmu

        

def formalSolnSph(numDeps, cosTheta, lambda2, tau, log10MinDepth, temp, lineMode,
                  radius, depths):
    """/**
 * Formal solution of the LTE radiative transfer for the monochromatic *surface*
 * intensity, I_lambda(Tau=0, theta) at wavelength lambda
 
  Spherical formal solution from Chapman, R., 1966, 143, 61  
  This solution is evaluated in geometric height coordinates wrt to centre of star (xi),
    NOT optical depth coordinates!
    
  # Kudritski Notes 3, Slide 26:
# Spherical rad trans eq:
# mu*dI/dr + {(1-mu^2)/r}*dI/dmu = eta - kappa*I
# So, DeltaI * {mu/DeltaZ + (1-mu^2)/(z*DeltaMu)} = eta - kappa * I
# - where mu(z) = mu at each z depth along a particualr ray of
#  initial mu=mu_0 from Gauss-Legendre quadrature  
 
 """

# *
# * Calls Planck.planck(lambda, temp) to get the LTE source function Input lambda
# * in nm for Planck
# */

    #logE = math.log10(math.e)  # // for debug output

    #Assume 32-bit limits to be on the safe side
    tiny = 1.175494e-38 
    tinySq = tiny**2
    huge = 3.4028237e38
    logTiny = math.log(tiny)
    #logHuge = math.log(huge)
    # // tolerance for stopping deeper contriibutions to I(Tau=0)
    #cutoff = 0.001

    # //  cosTheta is a 2xnumThetas array:
    # // row 0 is used for Gaussian quadrature weights
    # // row 1 is used for cos(theta) values
    # // Gaussian quadrature:
    # // Number of angles, numThetas, will have to be determined after the fact
    numThetas = len(cosTheta[0])

    
    #We'll make this very pedagogical by dressing everything up to look like 
    #Chapmsn (1966)'s notation:

    radiusCm = radius * Useful.rSun()
    #Take radius input parameter to be the radius of the bottom of the atmospheric model
    Ri = radiusCm
    logRi = math.log(Ri)
    #print("depths ", depths)
    
    ##Find depth index of optical depth unity:
    #t1Indx = min(range(len(tau[1])), key=lambda i: abs(tau[1][i]))
    ##print("t1Indx ", t1Indx, " tau[1] ", tau[1][t1Indx])
    #t1Dpth = depths[t1Indx] 
    
    #Find depth index of top of atmosphere 
    #t0Indx = min(range(len(tau[1])), key=lambda i: abs(tau[1][i] - log10MinDepth))
    #print("t0Indx ", t0Indx, " tau[1] ", tau[1][t0Indx])
    #topDpth = depths[t0Indx]     
    #Upper boundary condition affects first couple of depth values:
    t0Indx = 2
    
    thickness = depths[numDeps-1] - depths[t0Indx]
    thickness = thickness 
    
    #Take Chapman's R0 to be inner core radius + atmospheric thickness
    R0 = Ri + thickness
    #We're only interested in the surface intensity - fix Chapman's "r" variable to R0
    r = R0 # Surface intensity only
    #print("radiusCm ", radiusCm)
    
    #Construct Chapman's "xi" scale - the height coordinates in the atmosphere 
    # *relative to centre of star!
    #Height *decreases* with increasing index 
    testFudge = 1.0
    xi = [0.0 for i in range(numDeps)]
    xi[0] = R0 * testFudge

    #for i in range(0, numDeps):
    #    xi[i] = R0 - depths[i]
    xi = [(R0-z)*testFudge for z in depths]        
        
    #print("Ri ", Ri, " thickness ", thickness, " R0 ", R0)
    #print("")
    #print("xi ", xi)
        
    #For *surface* intensity, Chapman's critical angle wrt to the *surface*
    # normal of the I beam that grazes the inner core ("black sphere"):
    # CAUTION:  
    # I think this needs to be exactly consistent with R0 and Ri for
    # height integration purposes:
    #RiC = xi[t1Indx]
    R0C = xi[t0Indx]
    #No! muC = numpy.double(1.0) - math.pow(RiC/R0C, 2)
    muC = math.sqrt(1.0-math.pow(Ri/R0C, 2))
    if (muC < tinySq):
        muC = tiny
    else:
        muC = math.sqrt(muC)
    
    muC = abs(muC)
    #print("Ri ", Ri, " R0 ", R0, "thickness ", thickness, " muC ", muC)
    #test:
    #muC = 0.0001 #test
    
    
    #Chapman's integration factor function, capital Phi(r, mu):
    # At surface (r=R0):  Phi(R0, mu) = exp(0) = 1.0
    PhiR0 = numpy.double(1.0)
    logPhiR0 = math.log(PhiR0)
        
    #
    # //System.out.println("FORMALSOLN: numThetas= " + numThetas);
    # //double[][] intens = new double[numLams][numThetas];
    #intens = [numpy.double(0.0) for i in range(numThetas)]
    intensSph = [numpy.double(0.0) for i in range(numThetas)]  # Spherical solution
    intensPPz = [numpy.double(0.0) for i in range(numThetas)]  # For testing
    
    lineSourceVec = [numpy.double(0.0) for i in range(numDeps)]

    # //Get line source function vector, if needed:
    if (lineMode):
        lineSourceVec = LineProf.lineSource(numDeps, tau, temp, lambda2)

    if (lineMode == True):
        # //Line mode mode - ETLA + coherent scattering: S_lambda = (1-eps)*J_lambda + eps*B_lambda
        logSourceRi = lineSourceVec[numDeps-1]               
    else:
        # //Continuum mode - S_lambda = B_lambda
        logSourceRi = Planck.planck(temp[0][numDeps-1], lambda2)
    #SRi = math.exp(logSourceRi)
    
    #Get depth-dependent source function for all depths once and for all:
    logSXi = [numpy.double(0.0) for i in range(numDeps)]
    for id in range(0, numDeps):  
        if (lineMode == True):
            # //Line mode mode - ETLA + coherent scattering: S_lambda = (1-eps)*J_lambda + eps*B_lambda
            logSource = lineSourceVec[id]
        else:
            # //Continuum mode - S_lambda = B_lambda
            logSource = Planck.planck(temp[0][id], lambda2)
        logSXi[id] = logSource
        
    #We are only evaluating *surface* intensity at r = R0, so I *think*
    # the only mu values we need at the input Gauss-Legendre values for 
    #surface intensity

    # //for (int il = 0; il < numLams; il++ ) {
    for it in range(0, numThetas):
    #for it in range(0, 2):

        mu = cosTheta[1][it]
        
        #Factor (1-mu^2) recurs frequently:
        oneMinusMuSq = numpy.double(1.0) - math.pow(mu, 2)
        if (oneMinusMuSq < tiny):
            oneMinusMuSq = tiny
        
           
        if (mu >= muC):
            
            #print("it ", it, " mu ", mu)
            
            #muRi = 1.0 - ( ((r/Ri)**2) * oneMinusMuSq )
            logHelp = numpy.double(2.0)*(math.log(r)- logRi) + math.log(oneMinusMuSq)
            if (mu == muC):
                logHelp = -1.0 * logTiny
            muRi = numpy.double(1.0) - math.exp(logHelp)
            if (muRi < tinySq):
                muRi = tiny
            else:
                muRi = math.sqrt(muRi)
            #print("mu ", mu, " oneMinusMuSq ", oneMinusMuSq, " muRi ", muRi)
            PhiRi = Phi(numDeps-1, muRi, xi, t0Indx, it, tau)
            if (PhiRi < tiny):
                PhiRi = tiny
            logPhiRi = math.log(PhiRi)
            
            #Core intersecting beams:
            #term1 =  SRi * PhiRi / PhiR0
            logTerm1 = logSourceRi + logPhiRi - logPhiR0
            term1 = math.exp(logTerm1)
            #print("PhiRi ", PhiRi, " term1 ", term1)
            
            integrand = numpy.double(0.0)
            logIntegrand = numpy.double(0.0)
            integral = numpy.double(0.0)
            
            #height, xi, increases with decreasing i
            for id in range(numDeps-1, t0Indx, -1):  
            
                #muXi = 1.0 - ( ((r/xi[id])**2) * oneMinusMuSq )
                logHelp = numpy.double(2.0)*(math.log(r)-math.log(xi[id])) + math.log(oneMinusMuSq)
                muXi = numpy.double(1.0) - math.exp(logHelp)                
                if (muXi < tinySq):
                    muXi = tiny
                else:
                    muXi = math.sqrt(muXi)
                #print("id ", id, " muXi ", muXi, " logHelp ", logHelp, " r ", r,\
                #      " xi ", xi[id], " oneMinusMuSq ", oneMinusMuSq)
                #print("mu>MuC ", " id ", id, " mu ", mu, " muXi ", muXi)
                #print("logKappa ", logKappa)
                #print("rho[1] ", rho[1])
                PhiXi = Phi(id, muXi, xi, t0Indx, it, tau)
                if (PhiXi < tiny):
                    PhiXi = tiny
                #if (it == 0 and id%10 == 0):
                #    print("id ", id, " PhiXi ", PhiXi)
                
                #kappa =math.exp(logKappa[id])
                deltaXi = xi[id-1] - xi[id]  #Should be +ve
                
                #print("mu>muC " , " id ", id, " muXi ", muXi, " PhiXi ", PhiXi, "deltaXi ", deltaXi)
                #integrand = (PhiXi/PhiR0) * SXi[id] * kappa*rho[0][id] / muXi
                #print("it ", it, " id ", id, " log(PhiXi) ", math.log(PhiXi), " logPhiR0 ", logPhiR0,\
                #      " logSXi[id] ", logSXi[id], " log(muXi) ", math.log(muXi))
                ##      " logKappa ", logKappa[id], " rho[1] ", rho[1][id])
                #logIntegrand = (math.log(PhiXi) - logPhiR0) +\
                #    logSXi[id] + logKappa[id] + rho[1][id] + math.log(deltaXi) - math.log(muXi) 
                logIntegrand = (math.log(PhiXi) - logPhiR0) +\
                    logSXi[id] + math.log(tau[0][id]-tau[0][id-1]) - math.log(muXi)                
                integrand = math.exp(logIntegrand) 
                
                integral = integral + integrand
                #print("id ", id, " integrand ", integrand, " integral ", integral, " deltaXi ", deltaXi)
            intensSph[it] = term1 + integral
            #print("mu >= muC:")
            #print("it ", it, " term1 ", term1, " intensSph ", intensSph[it])
        
        """
        #Test simple version of spherical solution for special case of surface  
        #  intensity for central ray of mu=1:    
            
        if it == 0:

            Imu1 = 0.0
            
            int1 = 0.0
            term1 = 0.0
            for z in range(numDeps-1, t0Indx, -1):
                deltaXi = xi[z-1] - xi[z]
                integ = math.exp(logKappa[z])*rho[0][z]*deltaXi
                int1 = int1 + integ
                #print("z ", z, " integ ", integ)
            #logTerm1 = logSourceRi - int1
            #term1 = math.exp(logTerm1)
            term1 = math.exp(logSourceRi) * math.exp(-1.0*int1)
        
            int2 = 0.0
            int3 = 0.0
            logInt = 0.0
            for z in range(numDeps-1, t0Indx+1, -1):
                int2 = 0.0
                for z2 in range(z, t0Indx, -1):
                    deltaXi2 = xi[z2-1] - xi[z2]
                    logInteg2 = logKappa[z2] + rho[1][z2] + math.log(deltaXi2)
                    integ2 = math.exp(logInteg2) 
                    #if(z2%10 == 0):
                    #    print("z2 ", z2, " logKappa ", logKappa[z2], " rho[1] ", rho[1][z2],\
                    #          " deltaXi2 ", deltaXi2, " integ2 ", integ2)
                    #integ2 = math.exp(logKappa[z2])*rho[0][z2]*deltaXi2
                    #print("integ2 ", integ2)
                    int2 = int2 + integ2
                #if(z%10 == 0):
                #    print("z ", z, " int2 ", int2, " exp(-int2) ", math.exp(-1.0*int2))
                deltaXi = xi[z-1] - xi[z]
                logInt = logSXi[z] - int2 + logKappa[z] + rho[1][z] + math.log(deltaXi)
                int3 = int3 + math.exp(logInt) 
                #integ = math.exp(-1.0*int2) * math.exp(logSXi[z]) * math.exp(logKappa[z]) * rho[0][z] * deltaXi  
                #print("integ ", integ)
                #int3 = int3 + integ
                
            #print("term1 ", term1, " int3 ", int3)                
            Imu1 = term1 + int3
            
            print("mu ", mu, "intensSph[0] ", intensSph[it], " Imu1 ", Imu1)
        """
        
        """
        #Test plane-parallel formal solution for surface I transformed to geometric coordinate, z:
        #I(r, mu) = Int_(z=Ri)^(z=R0) {s(z) * e^[rho*kappa*(z-R0)/mu] * (rho*kappa/mu)} deltaZ
        # In domain z < r
        
        integral = 0.0
        for z in range(numDeps-1, t0Indx+1, -1):
            
            #Note: exponent should be -ve :  xi[z] < R0
            #exponent = math.exp(logKappa[z]) * rho[0][z] * (xi[z]-xi[t0Indx]) / mu
            #logIntegrand = logSXi[z] + logKappa[z] + rho[1][z] + exponent
            #integrand = math.exp(logIntegrand)
            #deltaXi = xi[z-1] - xi[z]  #+ve
            #integral = integral + integrand * deltaXi
           
            #For consistency with PP FormalSoln.formalSoln(), replace any kappa*rho*DelatXi
            # products with tau differences
            #Note: exponent should be -ve:
            exponent = -1.0 * tau[0][z] / mu
            logIntegrand = logSXi[z] + exponent
            integrand = math.exp(logIntegrand)
            integral = integral + (integrand * (tau[0][z]-tau[0][z-1]))
            
            ##Checking:
            #thisTau = math.exp(logKappa[z]) * rho[0][z] * (xi[z]-xi[t0Indx])
            #print("thisTau ", thisTau, " tau ", tau[0][z])
            
        intensPPz[it] = integral / mu
        #print("mu ", mu, "intensSph[0] ", intensSph[it], " integral ", integral)
        """
            
        if (mu < muC):
            
            #print("it ", it, " mu ", mu)
            #Beams that go through the outer atmosphere:

            #We cannot really hand in a vlaue of mu=0 to the integration factor function 
            # Phi(r, mu) - it will blow up at the lower limit of the integration                
            muMin = cosTheta[1][numThetas-1] 
            oneMinusMuSqMin = numpy.double(1.0) - math.pow(muMin, 2)
            #Test:
            #oneMinusMuSqMin = oneMinusMuSq
            #if (oneMinusMuSqMin < tiny):
            #    oneMinusMuSqMin = tiny
                
            #First integrate in depth to get I^- intensity value for first term:
                
            #Find lower limit of I^- integral:
            rLow = r * math.sqrt(oneMinusMuSq)
            logRLow = math.log(rLow)

            rLowIndx = min(range(len(xi)), key=lambda i: abs(xi[i] - rLow))
            #print("rLow ", rLow, " rLowIndx ", rLowIndx)
            if rLowIndx < 1:
                rLowIndx = 1
            #print("rLow ", rLow, " rLowIndx ", rLowIndx, " xi[rLowIndx] ", xi[rLowIndx])
            #print("xi[rl-2:rl+2] ", xi[rLowIndx-2:rLowIndx+3])

            #print("mu<muC", " mu ", mu, " muXi ", muXi)
            #print("logKappa ", logKappa)
            #print("rho[1] ", rho[1])            
            PhiRLow = Phi(rLowIndx, muMin, xi, t0Indx, it, tau)
            if (PhiRLow < tiny):
                PhiRLow = tiny
            logPhiRLow = math.log(PhiRLow)

            #print("PhiRLow ", PhiRLow)

            integrand = numpy.double(0.0)
            logIntegrand = numpy.double(0.0)
            integral = numpy.double(0.0)
            for id in range(rLowIndx-1, t0Indx, -1):
                
                #print("id ", id)
                #print("Before check: tiny ", tiny, " muXi ", muXi)
                #muXi = 1.0 - ( ((rLow/xi[id])**2)  * oneMinusMuSqMin )
                logHelp = numpy.double(2.0)*(logRLow-math.log(xi[id])) + math.log(oneMinusMuSqMin)
                muXi = numpy.double(1.0) - math.exp(logHelp)
                #print("it ", it, " id ", id, " muXi ", muXi, " help ", math.exp(logHelp))
                if (muXi < tinySq):
                    muXi = tiny
                    #print("Tiny branch")
                else:
                    muXi = math.sqrt(muXi)
                    #print("sqrt() branch")
                #print("After check: tiny ", tiny, " muXi ", muXi)
                #print("mu<muC", " id ", id, " mu ", mu, " muXi ", muXi)
                #print("logKappa ", logKappa)
                #print("rho[1] ", rho[1])
                PhiXi = Phi(id, muXi, xi, t0Indx, it, tau)
                if (PhiXi < tiny):
                    PhiXi = tiny
                
                #kappa =math.exp(logKappa[id])
                deltaXi = xi[id-1] - xi[id]  #Should be +ve
                #print("mu<muC, ", " id ", id, " muXi ", muXi, " PhiXi ", PhiXi, "deltaXi ", deltaXi)

                #print("Before division: tiny ", tiny, " muXi ", muXi, " PhiXi ", PhiXi)
                #integrand = (PhiRLow/PhiXi) * SXi[id] * kappa*rho[0][id] / muXi
                #print("it ", it, " id ", id, " logPhiRLow ", logPhiRLow, " log(PhiXi) ", math.log(PhiXi),\
                #      " logSXi[id] ", logSXi[id], " log(muXi) ", math.log(muXi))
                ##      " logKappa ", logKappa[id], " rho[1] ", rho[1][id])
                #logIntegrand = (logPhiRLow - math.log(PhiXi)) +\
                #    logSXi[id] + logKappa[id] + rho[1][id] + math.log(deltaXi) - math.log(muXi) 
                logIntegrand = (logPhiRLow - math.log(PhiXi)) +\
                    logSXi[id] + math.log(tau[0][id]-tau[0][id-1]) - math.log(muXi) 
                integrand = math.exp(logIntegrand) 
                #print("integrand ", integrand, " deltaXi ", deltaXi)
                #print("Division survivied")
                integral = integral + integrand
                #print("id ", id, " integrand ", integrand, " integral ", integral, " deltaXi ", deltaXi)
            Iminus = integral
            #print("mu < muC:")
            #print("muC ", muC, " rLowIndx ", rLowIndx, " it ", it, " Iminus ", Iminus)
            
            
            #term1 = Iminus * PhiRLow / PhiR0
            if Iminus < tiny:
                Iminus = tiny
            logTerm1 = math.log(Iminus) + logPhiRLow - logPhiR0
            term1 = math.exp(logTerm1)
            #print("it ", it, " Iminus ", Iminus, " term1 ", term1)
            
            integrand = numpy.double(0.0)
            logIntegrand = numpy.double(0.0)
            integral = numpy.double(0.0)
            
            #height, xi, increases with decreasing i
            for id in range(rLowIndx-1, t0Indx, -1):  
            
                #muXi = 1.0 - ( ((r/xi[id])**2) * oneMinusMuSq )
                logHelp = numpy.double(2.0)*(math.log(r)-math.log(xi[id])) + math.log(oneMinusMuSq)
                muXi = numpy.double(1.0) - math.exp(logHelp)                                
                if (muXi < tinySq):
                    muXi = tiny
                else:
                    muXi = math.sqrt(muXi)
                    
                PhiXi = Phi(id, muXi, xi, t0Indx, it, tau)
                if (PhiXi < tiny):
                    PhiXi = tiny
                
                #kappa =math.exp(logKappa[id])
                deltaXi = xi[id-1] - xi[id]  #Should be +ve

                #integrand = (PhiXi/PhiR0) * SXi[id] * kappa*rho[0][id] / muXi
                #print("muXi ", muXi)
                #logIntegrand = (math.log(PhiXi) - logPhiR0) +\
                #    logSXi[id] + logKappa[id] + rho[1][id] + math.log(deltaXi) - math.log(muXi) 
                logIntegrand = (math.log(PhiXi) - logPhiR0) +\
                    logSXi[id] + math.log(tau[0][id]-tau[0][id-1]) - math.log(muXi) 
                integrand = math.exp(logIntegrand)     
                integral = integral + integrand
                #print("id ", id, " integrand ", integrand, " integral ", integral)
                
            intensSph[it] = term1 + integral
            #print("mu < muC:")
            #print("muC ", muC)
            #print("it ", it, " integral ", integral, " intensSph ", intensSph[it])            
            
        

    #plt.plot(cosTheta[1], intens)

    # }   //it - theta loop

    # //} // il - lambda loop
    
    
    return intensSph
    #Testing return intensPPz  #Testing

