"""
Triboelectric Charging for Snow Particles
Snow Dust Project – Navigation via Electrostatic Fields
SCIENTIFICALLY CONSISTENT IMPLEMENTATION
"""

import numpy as np
from dataclasses import dataclass


@dataclass
class TriboelectricResult:
    charge_density: float
    charge_per_particle: float
    electric_field: float
    charging_rate: float
    polarity: str


class SnowTriboelectric:

    def __init__(self):
        # Empirical constants (order-of-magnitude realistic)
        self.k = 1.0e-15
        self.alpha = 1.2
        self.beta = 2.3
        self.gamma = 0.8

        # Work functions (eV)
        self.work_functions = {
            'ice': 3.7,
            'rime': 4.2,
            'graupel': 4.0,
            'snowflake': 3.9,
            'supercooled_water': 4.5,
            'dry_snow': 3.8,
            'wet_snow': 4.3
        }

        self.epsilon_0 = 8.854187817e-12

    # ==========================================================
    # SINGLE COLLISION CHARGE
    # ==========================================================
    def calculate_charge(self,
                         velocity: float,
                         diameter: float,
                         delta_temp: float,
                         material1: str = 'ice',
                         material2: str = 'snowflake',
                         humidity: float = 0.3,
                         pressure: float = 101325.0) -> TriboelectricResult:

        base_charge = self.k * \
                      (abs(velocity) ** self.alpha) * \
                      (diameter ** self.beta) * \
                      (abs(delta_temp) ** self.gamma)

        wf1 = self.work_functions.get(material1, 3.8)
        wf2 = self.work_functions.get(material2, 3.9)

        polarity = '-' if wf1 > wf2 else '+'
        polarity_factor = -1.0 if polarity == '-' else 1.0

        delta_wf = abs(wf1 - wf2)
        wf_factor = 1.0 + 0.3 * delta_wf

        humidity_factor = np.exp(-3.0 * humidity)
        pressure_factor = (101325.0 / pressure) ** 0.25

        final_charge = base_charge * polarity_factor * wf_factor * humidity_factor * pressure_factor

        volume = (np.pi / 6.0) * diameter**3
        charge_density = final_charge / volume if volume > 0 else 0.0

        e_field = abs(final_charge) / (4 * np.pi * self.epsilon_0)

        return TriboelectricResult(
            charge_density=charge_density,
            charge_per_particle=final_charge,
            electric_field=e_field,
            charging_rate=final_charge,
            polarity=polarity
        )

    # ==========================================================
    # PASCHEN BREAKDOWN FIELD
    # ==========================================================
    def breakdown_field(self, pressure: float, gap: float = 0.01) -> float:
        p_torr = pressure / 133.322
        d_cm = gap * 100.0
        pd = p_torr * d_cm

        if pd > 1:
            V_bd = (273.6 * pd) / np.log(6.89 * pd)
            return V_bd / gap
        else:
            return 3.0e6

    # ==========================================================
    # SNOWSTORM CHARGE BUILDUP MODEL
    # ==========================================================
    def simulate_snowstorm(self,
                           wind_speed: float,
                           mean_diameter: float,
                           temp_gradient: float,
                           duration: float,
                           particle_density: float):

        v_coll = 0.1 * wind_speed
        delta_T = temp_gradient * mean_diameter * 5.0

        single = self.calculate_charge(v_coll, mean_diameter, delta_T)

        sigma = np.pi * mean_diameter**2
        mean_free_path = 1.0 / (np.sqrt(2) * sigma * particle_density)
        f_coll = v_coll / mean_free_path if mean_free_path > 0 else 0.0

        total_collisions = f_coll * duration
        total_charge = single.charge_per_particle * total_collisions

        space_charge_density = total_charge * particle_density
        e_field_vertical = space_charge_density / (2 * self.epsilon_0)

        e_bd = self.breakdown_field(101325.0)
        risk = abs(e_field_vertical) / e_bd

        return {
            "collision_velocity": v_coll,
            "collision_frequency": f_coll,
            "charge_per_collision": single.charge_per_particle,
            "total_charge": total_charge,
            "space_charge_density": space_charge_density,
            "vertical_field": e_field_vertical,
            "breakdown_field": e_bd,
            "breakdown_risk": risk
        }


# ==============================================================
# TEST
# ==============================================================
if __name__ == "__main__":
    tribo = SnowTriboelectric()

    print("Testing triboelectric snow charging model\n")

    result = tribo.calculate_charge(
        velocity=2.5,
        diameter=200e-6,
        delta_temp=5.0,
        humidity=0.1
    )

    print(f"Single collision charge: {result.charge_per_particle:.2e} C")
    print(f"E-field at 1m: {result.electric_field:.1f} V/m")
    print(f"Polarity: {result.polarity}")

    storm = tribo.simulate_snowstorm(
        wind_speed=15.0,
        mean_diameter=300e-6,
        temp_gradient=0.5,
        duration=3600,
        particle_density=1e6
    )

    print("\nSnowstorm simulation:")
    print(f"Vertical E-field: {storm['vertical_field']:.1f} V/m")
    print(f"Breakdown risk: {storm['breakdown_risk']:.2%}")

    print("\nModel ready for electrostatic navigation research.")
