"""
Snowstorm Optical & Electrostatic Simulation
Combines Mie scattering (visibility) + Triboelectric charging
SciPy-free, Termux-compatible
"""

import numpy as np
from dataclasses import dataclass

# ==============================================================
# DATA STRUCTURES
# ==============================================================
@dataclass
class ScatteringResult:
    Q_sca: float
    Q_ext: float
    g: float
    sigma_ext: float

@dataclass
class TriboelectricResult:
    single_charge: float   # C
    E_field: float         # V/m at reference distance
    polarity: str
    breakdown_risk: float  # 0-1

# ==============================================================
# MIE SCATTERING MODEL
# ==============================================================
class SnowMieScattering:

    def __init__(self, wavelength=550e-9):
        self.wavelength = wavelength  # green light

    def compute(self, diameter, refractive_index=1.31):
        r = diameter / 2
        x = 2 * np.pi * r / self.wavelength

        if x < 1:
            return self._rayleigh(r, refractive_index)
        elif x < 2000:
            return self._mie_transition(r, x)
        else:
            return self._geometric(r)

    def _rayleigh(self, r, m):
        x = 2 * np.pi * r / self.wavelength
        Q_sca = (8/3) * x**4 * abs((m**2 - 1)/(m**2 + 2))**2
        Q_ext = Q_sca
        g = 0.0
        sigma_ext = Q_ext * np.pi * r**2
        return ScatteringResult(Q_sca, Q_ext, g, sigma_ext)

    def _mie_transition(self, r, x):
        Q_ray = (8/3) * x**4 * 0.01
        Q_geo = 2.0
        transition = 1 - np.exp(-0.12 * x)
        Q_sca = Q_ray * np.exp(-0.08 * x) + Q_geo * transition
        Q_ext = Q_sca
        g = 0.15 + 0.7 * (1 - np.exp(-x/8))
        sigma_ext = Q_ext * np.pi * r**2
        return ScatteringResult(Q_sca, Q_ext, g, sigma_ext)

    def _geometric(self, r):
        Q_sca = 2.0
        Q_ext = 2.0
        g = 0.87
        sigma_ext = Q_ext * np.pi * r**2
        return ScatteringResult(Q_sca, Q_ext, g, sigma_ext)

    def visibility(self, snow_density, diameter):
        result = self.compute(diameter)
        beta_ext = snow_density * result.sigma_ext
        if beta_ext > 0:
            return 3.912 / beta_ext
        else:
            return float('inf')

# ==============================================================
# TRIBOELECTRIC MODEL
# ==============================================================
class TriboelectricSnow:

    def __init__(self, particle_diameter=200e-6, density=1e6):
        self.diameter = particle_diameter
        self.density = density

    def compute_charge(self):
        # Approximate single collision charge (C)
        single_charge = 2.66e-23 * (self.diameter/200e-6)**1.5
        return single_charge

    def electric_field(self, distance=1.0):
        q = self.compute_charge() * self.density
        epsilon0 = 8.854e-12
        if distance > 0:
            E = q / (4 * np.pi * epsilon0 * distance**2)
        else:
            E = 0.0
        return E

    def breakdown_risk(self, distance=1.0):
        E = self.electric_field(distance)
        threshold = 3e6  # air breakdown V/m
        risk = min(E / threshold, 1.0)
        return risk

    def polarity(self):
        return "+" if np.random.rand() > 0.5 else "-"

    def simulate(self):
        q = self.compute_charge()
        E = self.electric_field()
        p = self.polarity()
        risk = self.breakdown_risk()
        return TriboelectricResult(q, E, p, risk)

# ==============================================================
# TEST SCRIPT
# ==============================================================
if __name__ == "__main__":
    print("=== Snowstorm Simulation ===\n")

    # Mie scattering test
    mie = SnowMieScattering()
    for d in [10e-6, 200e-6, 1e-3]:
        res = mie.compute(d)
        print(f"Particle {d*1e6:.0f} µm → Q_sca={res.Q_sca:.3f}, g={res.g:.2f}")

    vis = mie.visibility(1e6, 200e-6)
    print(f"Visibility in snowstorm: {vis:.1f} m\n")

    # Triboelectric test
    trib = TriboelectricSnow()
    trib_res = trib.simulate()
    print("Triboelectric snow simulation:")
    print(f"Single collision charge: {trib_res.single_charge:.2e} C")
    print(f"E-field at 1 m: {trib_res.E_field:.1f} V/m")
    print(f"Polarity: {trib_res.polarity}")
    print(f"Breakdown risk: {trib_res.breakdown_risk*100:.2f}%\n")

    print("Module ready for integrated snowstorm research.")
