import cv2
import numpy as np
import sys
from math import sqrt, acos, sin
from numpy import exp

"""
img = np.array(cv2.imread('noise.png', cv2.IMREAD_UNCHANGED))
img = img[:,:,:3]
"""
# CONSTANTS
airNumDensity       = 2.5035422e25
ozonePeak = 8e-6
ozoneNumDensity     = airNumDensity * exp(-35e3 / 8e3) * ozonePeak;
ozoneCrossSecR      = 4.51103766177301e-21 * 0.0001
ozoneCrossSecG      = 3.2854797958699e-21 * 0.0001
ozoneCrossSecB      = 1.96774621921165e-22 * 0.0001

rayleighCoeffR = 0.00000519673
rayleighCoeffG = 0.0000121427
rayleighCoeffB = 0.0000296453
mieCoeff = 8.6e-6
ozoneCoeffR    = ozoneCrossSecR*ozoneNumDensity
ozoneCoeffG    = ozoneCrossSecG*ozoneNumDensity
ozoneCoeffB    = ozoneCrossSecB*ozoneNumDensity

scaleHeightsRayl = 8500.0
scaleHeightsMie  = 1200.0

mieg = 0.77

planetRadius = 6371e3
atmosHeight  = 110e3
############

# HELPERS
def rsi(posX, posY, posZ, dirX, dirY, dirZ, r):
    b     = posX*dirX + posY*dirY + posZ*dirZ
    det   = b*b - (posX*posX + posY*posY + posZ*posZ) + r*r

    if det < 0.0:
        return -1.0

    det     = sqrt(det)

    return -b + det

def getRaylDensity(h):
    return exp(-h/scaleHeightsRayl)

def getMieDensity(h):
    return exp(-h/scaleHeightsMie)

def getOzoneDensity(elevation):
    ozonePeakAlt = 3e4
    ozoneFalloff = 6e3
    rOzoneFalloff = 1.0 / ozoneFalloff
    ozoneFalloffTop = 18e3
    rOzoneFalloffTop = 1.0 / ozoneFalloffTop
    d   = exp(-max(0.0, (ozonePeakAlt - elevation) - planetRadius) * rOzoneFalloff)
    d  *= exp(-max(0.0, (elevation - ozonePeakAlt) - planetRadius) * rOzoneFalloffTop)
    return d*0.7

def getElevation(rayPosX, rayPosY, rayPosZ):
    return sqrt(rayPosX*rayPosX + rayPosY*rayPosY + rayPosZ*rayPosZ) - planetRadius

##########
# optical density without extinction coefficients
def getAtmosPathDensity(rayPosX, rayPosY, rayPosZ, rayDirX, rayDirY, rayDirZ):
    dd = rsi(rayPosX, rayPosY, rayPosZ, rayDirX, rayDirY, rayDirZ, planetRadius + atmosHeight) / 128.0
    rayStepX = rayDirX*dd
    rayStepY = rayDirY*dd
    rayStepZ = rayDirZ*dd
    rayPosX = rayPosX + rayStepX*(0.5 * (max(rayDirY, 0.0) * 0.5 + 0.5))
    rayPosY = rayPosY + rayStepY*(0.5 * (max(rayDirY, 0.0) * 0.5 + 0.5))
    rayPosZ = rayPosZ + rayStepZ*(0.5 * (max(rayDirY, 0.0) * 0.5 + 0.5))

    odRayl  = 0.0
    odMie   = 0.0
    odOzone = 0.0
    for i in range(0, 128):
        elevation = getElevation(rayPosX, rayPosY, rayPosZ)

        odRayl += getRaylDensity(elevation) * dd
        odMie += getMieDensity(elevation) * dd
        odOzone += getOzoneDensity(elevation) * dd

        rayPosX += rayStepX
        rayPosY += rayStepY
        rayPosZ += rayStepZ
    return [odRayl / 10e3, odMie / 10e3, odOzone / 10e3]

# GENERATE LUT
img = np.zeros((128, 256, 3), dtype=np.float16)


for x in range(0, 256):
    for y in range(0, 128):
        cosTheta = (x/256.0)*2.0 - 1.0
        h = atmosHeight*y/128.0
    
        theta = acos(cosTheta)
        sinTheta = sin(theta)

        rayDirX = sinTheta
        rayDirY = cosTheta
        rayDirZ = 0.0
        od = getAtmosPathDensity(0.0, planetRadius + h, 0.0, rayDirX, rayDirY, rayDirZ)
        img[y,x] = od

img = img.astype('float16')
img.tofile('odLUT.dat')
print("finished generation")
print("data type: " + str(img.dtype))
print("size (h, w, channels): " + str(img.shape))
