defmodule Microwaveprop.Weather.SoundingParams do
  @moduledoc false

  @doc "Buck equation for saturation vapor pressure (hPa) given temperature in °C."
  @spec sat_vap_pres(float()) :: float()
  def sat_vap_pres(t_c) do
    6.1121 * :math.exp((18.678 - t_c / 234.5) * (t_c / (257.14 + t_c)))
  end

  @doc "Mixing ratio (g/kg) from dewpoint (°C) and pressure (hPa)."
  @spec mixing_ratio(float(), float()) :: float()
  def mixing_ratio(t_c, p_mb) do
    e = sat_vap_pres(t_c)
    622.0 * e / (p_mb - e)
  end

  @doc """
  Derive atmospheric propagation parameters from a sounding profile.

  Accepts either:
    * Legacy shape — string keys `"pres"` / `"hght"` / `"tmpc"` / `"dwpc"`
      (Elixir's HrrrClient, NARR, GEFS, UWYO clients).
    * Rust shape — `"pres_mb"` / `"hght_m"` / `"tmpc"` / `"dwpc"` (string
      or atom keys from the Rust hrrr_points worker and the f01..f18
      ProfilesFile cells).

  Both are normalized at entry to the canonical legacy shape so the
  rest of this module can assume `p["pres"]` / `p["hght"]` without per-
  source branching. Returns nil if fewer than 3 valid levels.
  """
  @spec derive([map()]) :: map() | nil
  def derive(profile) when is_list(profile) do
    sorted =
      profile
      |> Enum.map(&normalize_profile_entry/1)
      |> Enum.filter(fn p -> p["pres"] != nil and p["tmpc"] != nil and p["hght"] != nil end)
      |> Enum.sort_by(fn p -> p["pres"] end, :desc)

    if length(sorted) < 3 do
      nil
    else
      do_derive(sorted)
    end
  end

  @doc """
  Normalize a single profile entry to the canonical string-keyed shape
  (`"pres"` / `"hght"` / `"tmpc"` / `"dwpc"`). Tolerates:
    * The already-canonical shape (returned unchanged).
    * Rust's `"pres_mb"` / `"hght_m"` string keys.
    * Atom-keyed variants from `ProfilesFile.read/1` after Elixir's
      atomization pass.

  Exposed so `WeatherLayers` and other profile consumers share the
  same input contract.
  """
  @spec normalize_profile_entry(map()) :: map()
  def normalize_profile_entry(%{"pres" => _} = entry), do: entry

  def normalize_profile_entry(entry) when is_map(entry) do
    %{
      "pres" => pick(entry, ["pres_mb", :pres_mb, :pres]),
      "hght" => pick(entry, ["hght_m", :hght_m, "hght", :hght]),
      "tmpc" => pick(entry, ["tmpc", :tmpc]),
      "dwpc" => pick(entry, ["dwpc", :dwpc]),
      "drct" => pick(entry, ["drct", :drct]),
      "sknt" => pick(entry, ["sknt", :sknt])
    }
  end

  def normalize_profile_entry(entry), do: entry

  # Return the first non-nil value found under any of the candidate keys.
  defp pick(map, keys), do: Enum.find_value(keys, fn k -> Map.get(map, k) end)

  defp do_derive(sorted) do
    sfc = hd(sorted)
    sfc_hght = sfc["hght"]

    refract_profile = compute_refractivity_profile(sorted, sfc_hght)
    dn_dh = compute_gradients(refract_profile)
    inversions = detect_inversions(sorted, sfc_hght)
    ducts = detect_ducts(refract_profile)
    k_index = compute_k_index(sorted)
    li = compute_lifted_index(sorted, sfc)
    pw = compute_precipitable_water(sorted)
    bl_depth = compute_boundary_layer_depth(sorted, sfc)
    min_grad = find_min_gradient(dn_dh)

    sfc_n =
      case refract_profile do
        [first | _] -> first.n
        _ -> nil
      end

    %{
      level_count: length(sorted),
      surface_pressure_mb: sfc["pres"],
      surface_temp_c: sfc["tmpc"],
      surface_dewpoint_c: sfc["dwpc"] || sfc["tmpc"] - 10,
      surface_refractivity: sfc_n,
      min_refractivity_gradient: min_grad,
      boundary_layer_depth_m: bl_depth,
      precipitable_water_mm: pw,
      k_index: k_index,
      lifted_index: li,
      ducting_detected: ducts != [],
      ducts: ducts,
      duct_characteristics: if(ducts == [], do: nil, else: ducts),
      inversions: inversions,
      profile: sorted
    }
  end

  defp compute_refractivity_profile(sorted, sfc_hght) do
    sorted
    |> Enum.filter(fn p -> p["dwpc"] != nil end)
    |> Enum.map(fn p ->
      t_k = p["tmpc"] + 273.15
      e = sat_vap_pres(p["dwpc"])
      n = 77.6 * p["pres"] / t_k + 3.73e5 * e / (t_k * t_k)
      h_agl = p["hght"] - sfc_hght
      m = n + 0.157 * h_agl

      %{
        pres: p["pres"],
        h_agl: h_agl,
        h_km: h_agl / 1000.0,
        tmpc: p["tmpc"],
        dwpc: p["dwpc"],
        n: n,
        m: m
      }
    end)
  end

  defp compute_gradients(refract_profile) do
    refract_profile
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.flat_map(fn [prev, curr] ->
      dh_km = (curr.h_agl - prev.h_agl) / 1000.0

      # Only zero-thickness layers are skipped — the previous 10m guard
      # discarded the very thin near-surface levels that HRRR's native
      # vertical grid uses to resolve sharp inversions and surface ducts.
      if abs(dh_km) < 1.0e-6 do
        []
      else
        dn = curr.n - prev.n
        [%{h_agl: (curr.h_agl + prev.h_agl) / 2.0, gradient: dn / dh_km}]
      end
    end)
  end

  defp detect_inversions(sorted, sfc_hght) do
    raw_inversions =
      sorted
      |> Enum.chunk_every(2, 1, :discard)
      |> Enum.flat_map(fn [lower, upper] ->
        inversion_from_layer(lower, upper, sfc_hght)
      end)

    # Merge adjacent layers within 200m gap
    merged =
      raw_inversions
      |> Enum.reduce([], fn inv, acc -> merge_adjacent_inversion(inv, acc) end)
      |> Enum.reverse()

    # Keep only meaningful inversions: strength >= 0.5°C AND base below 5000m AGL
    Enum.filter(merged, fn inv -> inv.strength >= 0.5 and inv.base < 5000 end)
  end

  defp detect_ducts(refract_profile) do
    {ducts, in_duct_state} =
      refract_profile
      |> Enum.chunk_every(2, 1, :discard)
      |> Enum.reduce({[], nil}, fn [prev, curr], {ducts, duct_state} ->
        classify_duct_transition(prev, curr, ducts, duct_state)
      end)

    # Surface-based ducts that extend past the top of a shallow
    # profile still strongly affect ground-to-ground propagation, so
    # close the duct using the highest sampled level as the top
    # rather than discarding it for "lack of a clear top".
    ducts =
      case {in_duct_state, List.last(refract_profile)} do
        {%{base: base, base_m: base_m}, %{} = top} ->
          {finalized, _} = finalize_duct(ducts, base, base_m, top)
          finalized

        _ ->
          ducts
      end

    Enum.reverse(ducts)
  end

  defp inversion_from_layer(lower, upper, sfc_hght) do
    if upper["tmpc"] > lower["tmpc"] do
      base = lower["hght"] - sfc_hght
      top = upper["hght"] - sfc_hght
      strength = upper["tmpc"] - lower["tmpc"]
      dh_100m = (top - base) / 100.0
      rate = if dh_100m > 0, do: strength / dh_100m, else: 0.0
      [%{base: base, top: top, strength: strength, rate: rate}]
    else
      []
    end
  end

  defp merge_adjacent_inversion(inv, []) do
    [inv]
  end

  defp merge_adjacent_inversion(inv, [last | rest] = acc) do
    if inv.base <= last.top + 200 do
      merged_inv = %{last | top: inv.top, strength: last.strength + inv.strength}
      [merged_inv | rest]
    else
      [inv | acc]
    end
  end

  defp classify_duct_transition(prev, curr, ducts, duct_state) do
    dm = curr.m - prev.m

    case {dm < 0, duct_state} do
      {true, nil} ->
        # Entering a duct
        {ducts, %{base: prev.h_agl, base_m: prev.m}}

      {false, %{base: base, base_m: base_m}} ->
        # Exiting a duct
        finalize_duct(ducts, base, base_m, prev)

      _ ->
        {ducts, duct_state}
    end
  end

  defp finalize_duct(ducts, base, base_m, prev) do
    top = prev.h_agl
    strength = base_m - prev.m

    if strength > 2 do
      duct = %{"base" => base, "top" => top, "strength" => Float.round(strength, 1)}
      {[duct | ducts], nil}
    else
      {ducts, nil}
    end
  end

  defp compute_k_index(sorted) do
    p850 = nearest_level(sorted, 850)
    p700 = nearest_level(sorted, 700)
    p500 = nearest_level(sorted, 500)

    with %{"tmpc" => t850} when t850 != nil <- p850,
         %{"tmpc" => t700} when t700 != nil <- p700,
         %{"tmpc" => t500} when t500 != nil <- p500 do
      td850 = p850["dwpc"] || -30.0
      td700 = p700["dwpc"] || -30.0
      t850 - t500 + td850 - (t700 - td700)
    else
      _ -> nil
    end
  end

  defp compute_lifted_index(sorted, sfc) do
    p500 = nearest_level(sorted, 500)

    case p500 do
      %{"tmpc" => t500, "hght" => h500} when t500 != nil and h500 != nil ->
        sfc_tmpc = sfc["tmpc"]
        sfc_hght = sfc["hght"]
        # Simplified LI: T500 - (Tsfc - lapse_rate * height_diff)
        t500 - (sfc_tmpc - (h500 - sfc_hght) * 0.00976)

      _ ->
        nil
    end
  end

  defp compute_precipitable_water(sorted) do
    sorted
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.reduce(0.0, fn [lower, upper], pw ->
      if lower["dwpc"] != nil and upper["dwpc"] != nil do
        mr1 = mixing_ratio(lower["dwpc"], lower["pres"])
        mr2 = mixing_ratio(upper["dwpc"], upper["pres"])
        dp = lower["pres"] - upper["pres"]
        pw + (mr1 + mr2) / 2.0 * dp / 9.81 / 10.0
      else
        pw
      end
    end)
  end

  defp compute_boundary_layer_depth(sorted, sfc) do
    sfc_hght = sfc["hght"]
    sfc_tmpc = sfc["tmpc"]
    theta_sfc = sfc_tmpc + 273.15 + 9.8 * (sfc_hght / 1000.0)

    sorted
    |> tl()
    |> Enum.find_value(fn p ->
      theta = p["tmpc"] + 273.15 + 9.8 * (p["hght"] / 1000.0)

      if theta > theta_sfc + 2 do
        p["hght"] - sfc_hght
      end
    end)
  end

  defp find_min_gradient([]), do: nil

  defp find_min_gradient(dn_dh) do
    dn_dh
    |> Enum.min_by(fn %{gradient: g} -> g end)
    |> Map.get(:gradient)
  end

  defp nearest_level(sorted, target_pres) do
    Enum.min_by(sorted, fn p -> abs(p["pres"] - target_pres) end)
  end
end