suns 0.0.1

use sciforge::hub::domain::common::constants::{K_B, SIGMA_SB};
use sciforge::hub::domain::meteorology::radiation::stefan_boltzmann_flux;

use crate::{
    PROTON_MASS, SOLAR_CORE_DENSITY, SOLAR_CORE_TEMP, SOLAR_EFFECTIVE_TEMP,
    SOLAR_RADIUS, SPEED_OF_LIGHT,
};

pub const RADIATIVE_OPACITY_CORE: f64 = 0.4;
pub const RADIATIVE_OPACITY_ENVELOPE: f64 = 1.0;
pub const THOMSON_CROSS_SECTION: f64 = 6.6524e-29;
pub const PHOTON_MEAN_FREE_PATH_CORE: f64 = 0.01;
pub const PHOTON_DIFFUSION_TIMESCALE_YR: f64 = 1.7e5;
pub const CONVECTIVE_VELOCITY_SURFACE: f64 = 1000.0;
pub const MIXING_LENGTH_PARAMETER: f64 = 1.7;
pub const GRANULATION_SIZE_M: f64 = 1e6;
pub const SUPERGRANULATION_SIZE_M: f64 = 3e7;
pub const GRANULE_LIFETIME_S: f64 = 600.0;

pub struct RadiativeTransfer {
    pub temperature_k: f64,
    pub density_kgm3: f64,
    pub opacity_cm2g: f64,
    pub radius_m: f64,
}

impl RadiativeTransfer {
    pub fn mean_free_path(&self) -> f64 {
        1.0 / (self.opacity_cm2g * 0.1 * self.density_kgm3)
    }

    pub fn optical_depth(&self, path_length: f64) -> f64 {
        self.opacity_cm2g * 0.1 * self.density_kgm3 * path_length
    }

    pub fn radiative_flux(&self) -> f64 {
        stefan_boltzmann_flux(self.temperature_k)
    }

    pub fn radiation_pressure(&self) -> f64 {
        4.0 * SIGMA_SB * self.temperature_k.powi(4) / (3.0 * SPEED_OF_LIGHT)
    }

    pub fn luminosity_at_radius(&self) -> f64 {
        4.0 * std::f64::consts::PI * self.radius_m.powi(2) * self.radiative_flux()
    }

    pub fn temperature_gradient_radiative(&self) -> f64 {
        let kappa = self.opacity_cm2g * 0.1;
        let l_r = self.luminosity_at_radius();
        3.0 * kappa * self.density_kgm3 * l_r
            / (16.0
                * std::f64::consts::PI
                * SPEED_OF_LIGHT
                * 4.0
                * SIGMA_SB
                * self.temperature_k.powi(3)
                * self.radius_m.powi(2))
    }

    pub fn photon_diffusion_speed(&self) -> f64 {
        SPEED_OF_LIGHT * self.mean_free_path() / (3.0 * self.radius_m)
    }

    pub fn rosseland_mean_opacity(&self) -> f64 {
        self.opacity_cm2g
    }
}

pub fn core_radiative_transfer() -> RadiativeTransfer {
    RadiativeTransfer {
        temperature_k: SOLAR_CORE_TEMP,
        density_kgm3: SOLAR_CORE_DENSITY,
        opacity_cm2g: RADIATIVE_OPACITY_CORE,
        radius_m: 0.25 * SOLAR_RADIUS,
    }
}

pub fn radiative_zone_transfer() -> RadiativeTransfer {
    RadiativeTransfer {
        temperature_k: 4.0e6,
        density_kgm3: 1e3,
        opacity_cm2g: RADIATIVE_OPACITY_ENVELOPE,
        radius_m: 0.5 * SOLAR_RADIUS,
    }
}

pub struct ConvectiveTransport {
    pub temperature_k: f64,
    pub density_kgm3: f64,
    pub pressure_scale_height: f64,
    pub gravity: f64,
    pub delta_t: f64,
}

impl ConvectiveTransport {
    pub fn mixing_length(&self) -> f64 {
        MIXING_LENGTH_PARAMETER * self.pressure_scale_height
    }

    pub fn convective_velocity(&self) -> f64 {
        let l = self.mixing_length();
        let buoyancy = self.gravity * self.delta_t / self.temperature_k;
        (buoyancy * l).sqrt()
    }

    pub fn convective_flux(&self) -> f64 {
        let v = self.convective_velocity();
        let cp = 5.0 * K_B / (2.0 * PROTON_MASS);
        self.density_kgm3 * cp * v * self.delta_t
    }

    pub fn convective_luminosity(&self, radius_m: f64) -> f64 {
        4.0 * std::f64::consts::PI * radius_m.powi(2) * self.convective_flux()
    }

    pub fn convective_turnover_time(&self) -> f64 {
        self.mixing_length() / self.convective_velocity().max(1e-10)
    }

    pub fn brunt_vaisala_frequency(&self) -> f64 {
        let g = self.gravity;
        let gamma = 5.0 / 3.0;
        let grad_ad = 1.0 - 1.0 / gamma;
        let grad_actual = self.delta_t / self.temperature_k;
        let n2 = g / self.pressure_scale_height * (grad_actual - grad_ad);
        if n2 > 0.0 { n2.sqrt() } else { 0.0 }
    }
}

pub fn surface_convection() -> ConvectiveTransport {
    ConvectiveTransport {
        temperature_k: SOLAR_EFFECTIVE_TEMP,
        density_kgm3: 2e-4,
        pressure_scale_height: 1.5e5,
        gravity: 274.0,
        delta_t: 500.0,
    }
}

pub fn deep_convection() -> ConvectiveTransport {
    ConvectiveTransport {
        temperature_k: 2.0e6,
        density_kgm3: 200.0,
        pressure_scale_height: 5e7,
        gravity: 500.0,
        delta_t: 1e4,
    }
}

pub fn photon_random_walk_time() -> f64 {
    let n_steps = (SOLAR_RADIUS / PHOTON_MEAN_FREE_PATH_CORE).powi(2);
    n_steps * PHOTON_MEAN_FREE_PATH_CORE / SPEED_OF_LIGHT
}

pub fn schwarzschild_criterion(grad_rad: f64, grad_ad: f64) -> bool {
    grad_rad > grad_ad
}

pub fn adiabatic_gradient() -> f64 {
    1.0 - 1.0 / (5.0 / 3.0)
}

pub fn kramers_opacity(density: f64, temperature: f64) -> f64 {
    let kappa_0 = 4.34e25;
    kappa_0 * density * temperature.powf(-3.5)
}

pub fn electron_scattering_opacity(hydrogen_fraction: f64) -> f64 {
    0.2 * (1.0 + hydrogen_fraction)
}

pub fn eddington_flux(temperature: f64) -> f64 {
    SIGMA_SB * temperature.powi(4) / std::f64::consts::PI
}

pub fn radiative_equilibrium_temperature(luminosity: f64, radius: f64) -> f64 {
    (luminosity / (4.0 * std::f64::consts::PI * radius * radius * SIGMA_SB)).powf(0.25)
}

pub fn limb_darkening(theta: f64) -> f64 {
    let u = 0.6;
    1.0 - u * (1.0 - theta.cos())
}

pub fn granulation_power_spectrum(wavenumber: f64) -> f64 {
    let k0 = 2.0 * std::f64::consts::PI / GRANULATION_SIZE_M;
    let v0 = CONVECTIVE_VELOCITY_SURFACE;
    v0 * v0 / (1.0 + (wavenumber / k0).powi(4))
}