defmodule Microwaveprop.Weather.SkewtParams do
  @moduledoc """
  Derived sounding parameters for the `/skewt` page — the same set the
  SPC sounding viewer publishes in its `*.txt` companion file
  (`https://www.spc.noaa.gov/exper/soundings/`):

    * Lifting condensation level (`lcl_pressure_mb`, `lcl_temp_c`,
      `lcl_height_m`).
    * Level of free convection (`lfc_pressure_mb`, `lfc_height_m`).
    * Equilibrium level (`el_pressure_mb`, `el_height_m`).
    * Surface-based CAPE/CIN (`sbcape`, `sbcin` in J/kg).
    * Three-kilometre CAPE (`cape_3km`).
    * Lifted index (`lifted_index`, parcel-environment T diff at 500 mb).
    * Freezing level and wet-bulb-zero level (`freezing_level_m`,
      `wbz_m`).
    * Downdraft CAPE (`dcape`).
    * Layer lapse rates (`lapse_rate_sfc_3km_c_per_km`,
      `lapse_rate_700_500_c_per_km`, `lapse_rate_850_500_c_per_km`).
    * `parcel_trace` — the surface parcel's pressure/temperature
      trajectory through the diagram, used by `SkewtSvg.render/1` to
      draw the dashed parcel-ascent overlay.

  Wind-derived indices (SRH, BWD, Bunkers motion, SCP, STP, SHIP) are
  *not* produced because HRRR persists wind only at 10 m AGL — there
  is no per-pressure-level wind to compute them from. Operators wanting
  shear-based products should use NUCAPS / SPC mesoanalysis instead.

  Inputs are the same canonical profile shape `SoundingParams.derive/1`
  consumes: `[%{"pres", "hght", "tmpc", "dwpc"}, ...]` (string OR atom
  keys; a `:hght_m` / `"hght_m"` alias for height is also accepted).
  """

  @r_d 287.05
  @r_v 461.5
  @c_p 1004.0
  @g 9.80665
  @epsilon @r_d / @r_v
  # Buck's saturation-vapor formula constants (kept in this module so
  # we don't have to leak them out of SoundingParams).
  @buck_a 6.1121
  @buck_b 17.502
  @buck_c 240.97
  # Approximate L_v at 0 °C in J/kg. Tropospheric variation across
  # parcel temperatures is ≤4 % which is well below CAPE bin noise.
  @l_v 2.5e6

  @nil_params %{
    lcl_pressure_mb: nil,
    lcl_temp_c: nil,
    lcl_height_m: nil,
    lfc_pressure_mb: nil,
    lfc_height_m: nil,
    el_pressure_mb: nil,
    el_height_m: nil,
    sbcape: nil,
    sbcin: nil,
    cape_3km: nil,
    lifted_index: nil,
    freezing_level_m: nil,
    wbz_m: nil,
    dcape: nil,
    lapse_rate_sfc_3km_c_per_km: nil,
    lapse_rate_700_500_c_per_km: nil,
    lapse_rate_850_500_c_per_km: nil,
    parcel_trace: []
  }

  @doc """
  Returns the SPC-style parameter map for `profile`. See module
  docstring for field-by-field semantics. Returns `@nil_params` for
  empty / single-level inputs.
  """
  @spec derive([map()]) :: map()
  def derive(profile) when is_list(profile) do
    levels = profile |> Enum.map(&normalize_level/1) |> Enum.reject(&is_nil/1)

    case Enum.sort_by(levels, & &1.pres, :desc) do
      [] -> @nil_params
      [_only] -> @nil_params
      [surface | _] = sorted -> compute(surface, sorted)
    end
  end

  defp compute(surface, sorted_desc) do
    {p_lcl, t_lcl} = lcl(surface.tmpc, surface.dwpc, surface.pres)
    parcel_trace = parcel_trace(surface, p_lcl, t_lcl, sorted_desc)
    {sbcape, sbcin} = cape_cin(parcel_trace, sorted_desc)
    cape_3km = cape_layer(parcel_trace, sorted_desc, surface.hght, surface.hght + 3000.0)
    {lfc_pressure, el_pressure} = lfc_el(parcel_trace, sorted_desc)

    {lfc_height, el_height} =
      {height_at_pressure(sorted_desc, lfc_pressure), height_at_pressure(sorted_desc, el_pressure)}

    %{
      lcl_pressure_mb: p_lcl,
      lcl_temp_c: t_lcl,
      lcl_height_m: height_at_pressure(sorted_desc, p_lcl),
      lfc_pressure_mb: lfc_pressure,
      lfc_height_m: lfc_height,
      el_pressure_mb: el_pressure,
      el_height_m: el_height,
      sbcape: sbcape,
      sbcin: sbcin,
      cape_3km: cape_3km,
      lifted_index: lifted_index(parcel_trace, sorted_desc),
      freezing_level_m: temperature_height_crossing(sorted_desc, & &1.tmpc, 0.0),
      wbz_m: temperature_height_crossing(sorted_desc, &wet_bulb_c/1, 0.0),
      dcape: dcape(sorted_desc),
      lapse_rate_sfc_3km_c_per_km: lapse_rate_height(sorted_desc, surface.hght, surface.hght + 3000.0),
      lapse_rate_700_500_c_per_km: lapse_rate_pressures(sorted_desc, 700.0, 500.0),
      lapse_rate_850_500_c_per_km: lapse_rate_pressures(sorted_desc, 850.0, 500.0),
      parcel_trace: parcel_trace
    }
  end

  # ── Lifting Condensation Level (Bolton 1980 eq 22) ──────────────────

  @doc """
  Lifting condensation level pressure and temperature.

  Inputs:
    * `temp_c` — surface temperature, °C
    * `dewpoint_c` — surface dewpoint, °C
    * `pressure_mb` — surface pressure, mb

  Returns `{p_lcl_mb, t_lcl_c}`. Saturated air (T = T_d) returns the
  surface pressure unchanged.
  """
  @spec lcl(number(), number(), number()) :: {float(), float()}
  def lcl(temp_c, dewpoint_c, pressure_mb) when is_number(temp_c) and is_number(dewpoint_c) and is_number(pressure_mb) do
    t_k = temp_c + 273.15
    td_k = dewpoint_c + 273.15

    # Bolton (1980) eq 22: T at LCL.
    t_lcl_k = 1.0 / (1.0 / (td_k - 56.0) + :math.log(t_k / td_k) / 800.0) + 56.0

    # Poisson relation between two adiabatic levels.
    p_lcl = pressure_mb * :math.pow(t_lcl_k / t_k, @c_p / @r_d)

    {p_lcl, t_lcl_k - 273.15}
  end

  # ── Parcel ascent ───────────────────────────────────────────────────

  # Returns a list of `%{pres, tmpc}` for the surface parcel, walking
  # from surface pressure up through every environmental level above
  # the LCL using a saturated-adiabatic step in pressure.
  defp parcel_trace(surface, p_lcl, t_lcl, sorted_desc) do
    surface_step = %{pres: surface.pres, tmpc: surface.tmpc}
    lcl_step = %{pres: min(p_lcl, surface.pres), tmpc: t_lcl}

    above_lcl =
      sorted_desc
      |> Enum.filter(&(&1.pres < p_lcl))
      |> Enum.sort_by(& &1.pres, :desc)

    {trace_above, _} =
      Enum.map_reduce(above_lcl, lcl_step, fn level, prev ->
        t_next = saturated_adiabatic_step(prev.tmpc, prev.pres, level.pres)
        step = %{pres: level.pres, tmpc: t_next}
        {step, step}
      end)

    Enum.uniq_by([surface_step, lcl_step | trace_above], & &1.pres)
  end

  # Integrate dT/d(ln p) = γ_sat × R_d × T / g across the pressure
  # interval [p1, p2] in 5 mb steps using a forward-Euler scheme.
  # Coarse but stable; the residual error is well below CAPE-bin noise.
  defp saturated_adiabatic_step(t_c_start, p_start, p_end) do
    n_steps = max(1, trunc((p_start - p_end) / 5.0))
    dp = (p_end - p_start) / n_steps

    1..n_steps
    |> Enum.reduce({t_c_start, p_start}, fn _, {t_c, p} ->
      t_k = t_c + 273.15
      gamma_s = saturated_lapse_rate_per_pressure(t_k, p)
      t_next_k = t_k + gamma_s * dp
      {t_next_k - 273.15, p + dp}
    end)
    |> elem(0)
  end

  # γ_sat expressed as dT/dp (K/mb), derived from the standard moist
  # adiabatic lapse rate dT/dz combined with hydrostatic dz/dp =
  # -R_d × T_v / (g × p).
  defp saturated_lapse_rate_per_pressure(t_k, p_mb) do
    e_s = saturation_vapor_pressure_mb(t_k - 273.15)
    r_s = @epsilon * e_s / max(p_mb - e_s, 1.0e-3)

    num = 1.0 + @l_v * r_s / (@r_d * t_k)
    den = @c_p + @l_v * @l_v * r_s * @epsilon / (@r_d * t_k * t_k)
    gamma_z = @g * num / den

    # dz/dp ≈ -R_d × T / (g × p)  (mb)
    dz_dp = -@r_d * t_k / (@g * p_mb)
    -gamma_z * dz_dp
  end

  defp saturation_vapor_pressure_mb(t_c) do
    @buck_a * :math.exp(@buck_b * t_c / (@buck_c + t_c))
  end

  # ── CAPE / CIN ──────────────────────────────────────────────────────

  defp cape_cin(parcel_trace, sorted_desc) do
    pairs = parcel_environment_pairs(parcel_trace, sorted_desc)

    {cape, cin} =
      Enum.reduce(layer_increments(pairs), {0.0, 0.0}, fn {avg_diff, dlnp}, {cape, cin} ->
        delta = @r_d * avg_diff * dlnp

        cond do
          avg_diff > 0 -> {cape + delta, cin}
          avg_diff < 0 -> {cape, cin + delta}
          true -> {cape, cin}
        end
      end)

    {round_j(cape), round_j(cin)}
  end

  defp cape_layer(_parcel, _profile, lo, hi) when not (is_number(lo) and is_number(hi)), do: nil

  defp cape_layer(parcel_trace, sorted_desc, lo_height_m, hi_height_m) do
    pairs =
      parcel_trace
      |> parcel_environment_pairs(sorted_desc)
      |> Enum.filter(fn p -> p.height >= lo_height_m and p.height <= hi_height_m end)

    cape =
      pairs
      |> layer_increments()
      |> Enum.reduce(0.0, fn {avg_diff, dlnp}, acc ->
        if avg_diff > 0, do: acc + @r_d * avg_diff * dlnp, else: acc
      end)

    round_j(cape)
  end

  # Build a list of {pressure, parcel_T_K, env_T_K, height} pairs at
  # every environmental pressure level the parcel trace also covers.
  defp parcel_environment_pairs(parcel_trace, sorted_desc) do
    parcel_by_pres = Map.new(parcel_trace, &{&1.pres, &1.tmpc})

    Enum.flat_map(sorted_desc, fn level ->
      case Map.fetch(parcel_by_pres, level.pres) do
        {:ok, parcel_t} ->
          [
            %{
              pres: level.pres,
              parcel_t_k: parcel_t + 273.15,
              env_t_k: level.tmpc + 273.15,
              diff: parcel_t - level.tmpc,
              height: level.hght
            }
          ]

        :error ->
          []
      end
    end)
  end

  defp layer_increments(pairs) do
    pairs
    |> Enum.sort_by(& &1.pres, :desc)
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.map(fn [a, b] ->
      avg = (a.diff + b.diff) / 2.0
      dlnp = :math.log(a.pres / b.pres)
      {avg, dlnp}
    end)
  end

  # ── LFC, EL ─────────────────────────────────────────────────────────

  defp lfc_el(parcel_trace, sorted_desc) do
    pairs =
      parcel_trace
      |> parcel_environment_pairs(sorted_desc)
      |> Enum.sort_by(& &1.pres, :desc)

    pos_pairs = Enum.filter(pairs, &(&1.diff > 0))

    case pos_pairs do
      [] ->
        {nil, nil}

      [first | _] ->
        last = List.last(pos_pairs)
        # LFC = bottom of the first contiguous positive-area region.
        # EL = top of the last positive-area region. Coarse but
        # matches what SPC publishes for single-parcel ascent.
        {first.pres, last.pres}
    end
  end

  # ── Lifted index ────────────────────────────────────────────────────

  defp lifted_index(parcel_trace, sorted_desc) do
    parcel_t = pressure_lookup_temp(parcel_trace, 500.0)
    env_t = pressure_lookup_temp(sorted_desc, 500.0)

    case {parcel_t, env_t} do
      {nil, _} -> nil
      {_, nil} -> nil
      {pt, et} -> Float.round(et - pt, 1)
    end
  end

  defp pressure_lookup_temp(levels, target_p) do
    sorted = Enum.sort_by(levels, & &1.pres, :desc)

    cond do
      sorted == [] -> nil
      target_p > List.first(sorted).pres -> List.first(sorted).tmpc
      target_p < List.last(sorted).pres -> List.last(sorted).tmpc
      true -> interp_temp_at_pressure(sorted, target_p)
    end
  end

  defp interp_temp_at_pressure(sorted, target_p) do
    sorted
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.find_value(fn [a, b] ->
      if target_p <= a.pres and target_p >= b.pres do
        f = (a.pres - target_p) / max(a.pres - b.pres, 1.0e-6)
        a.tmpc + f * (b.tmpc - a.tmpc)
      end
    end)
  end

  # ── Freezing / wet-bulb-zero / DCAPE ────────────────────────────────

  defp temperature_height_crossing(sorted_desc, fun, target_c) do
    sorted_desc
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.find_value(fn [a, b] ->
      ta = fun.(a)
      tb = fun.(b)

      cond do
        ta == nil or tb == nil ->
          nil

        (ta - target_c) * (tb - target_c) > 0 ->
          nil

        ta == tb ->
          a.hght

        true ->
          f = (ta - target_c) / (ta - tb)
          a.hght + f * (b.hght - a.hght)
      end
    end)
  end

  @doc """
  Wet-bulb temperature (°C) for a level using a Stipanuk-style
  iterative bracket between dewpoint and dry bulb. Returns `nil` if
  either temperature is missing.

  Accepts a level map with `tmpc`, `dwpc`, and `pres` keys (the same
  normalized shape `derive/1` consumes).
  """
  @spec wet_bulb_c(map()) :: float() | nil
  def wet_bulb_c(%{tmpc: t, dwpc: t_d, pres: _p}) when t == nil or t_d == nil, do: nil

  def wet_bulb_c(%{tmpc: t, dwpc: t_d, pres: p}) do
    bisect_wet_bulb(t, t_d, p)
  end

  def wet_bulb_c(_), do: nil

  @doc """
  Mixing ratio (kg/kg) at a given dewpoint and pressure, via the Buck
  saturation-vapor formula evaluated at the dewpoint temperature.
  Returns `nil` for non-numeric inputs.
  """
  @spec mixing_ratio(number() | nil, number() | nil) :: float() | nil
  def mixing_ratio(dewpoint_c, pressure_mb) when is_number(dewpoint_c) and is_number(pressure_mb) do
    e = saturation_vapor_pressure_mb(dewpoint_c)
    @epsilon * e / max(pressure_mb - e, 1.0e-3)
  end

  def mixing_ratio(_, _), do: nil

  @doc """
  Virtual temperature (°C) for a parcel/level. Adds the moisture
  correction `(1 + 0.61 r)` to the dry-bulb temperature, where `r` is
  the mixing ratio in kg/kg.

  Two calling forms:

    * `virtual_temp_c(temp_c, dewpoint_c, pressure_mb)` — environment
      virtual temperature (uses the actual mixing ratio from the
      dewpoint).
    * `virtual_temp_c(temp_c, pressure_mb)` — parcel virtual temperature
      assuming saturation (uses the saturation mixing ratio at the
      parcel's own temperature). This is what CAPE uses above the LCL.
  """
  @spec virtual_temp_c(number(), number(), number()) :: float()
  def virtual_temp_c(temp_c, dewpoint_c, pressure_mb) do
    r = mixing_ratio(dewpoint_c, pressure_mb) || 0.0
    (temp_c + 273.15) * (1.0 + 0.61 * r) - 273.15
  end

  @spec virtual_temp_c(number(), number()) :: float()
  def virtual_temp_c(temp_c, pressure_mb) do
    virtual_temp_c(temp_c, temp_c, pressure_mb)
  end

  defp bisect_wet_bulb(t, t_d, _p) when t == t_d, do: t

  defp bisect_wet_bulb(t, t_d, p) do
    e = saturation_vapor_pressure_mb(t_d)

    f = fn t_w ->
      e_s = saturation_vapor_pressure_mb(t_w)
      e_s - e - p * @c_p / (@epsilon * @l_v) * (t - t_w)
    end

    bisect(f, t_d, t, 30)
  end

  defp bisect(_f, lo, hi, 0), do: (lo + hi) / 2.0

  defp bisect(f, lo, hi, iter) do
    mid = (lo + hi) / 2.0
    val_lo = f.(lo)
    val_mid = f.(mid)

    cond do
      abs(val_mid) < 1.0e-3 -> mid
      val_lo * val_mid < 0 -> bisect(f, lo, mid, iter - 1)
      true -> bisect(f, mid, hi, iter - 1)
    end
  end

  defp dcape(sorted_desc) do
    # Pick the level with minimum θ_e in the layer 700-400 mb as the
    # downdraft source. Saturate the parcel there, descend moist-
    # adiabatically to the surface, and integrate negative buoyancy.
    candidates =
      Enum.filter(sorted_desc, fn l -> l.pres >= 400.0 and l.pres <= 700.0 end)

    case candidates do
      [] ->
        nil

      _ ->
        source = Enum.min_by(candidates, &theta_e(&1.tmpc, &1.dwpc, &1.pres))
        descend_dcape(source, sorted_desc)
    end
  end

  defp descend_dcape(source, sorted_desc) do
    descent_targets =
      sorted_desc
      |> Enum.filter(&(&1.pres >= source.pres))
      |> Enum.sort_by(& &1.pres, :asc)

    {traj, _} =
      Enum.map_reduce(descent_targets, %{pres: source.pres, tmpc: source.dwpc}, fn lvl, prev ->
        t_next = saturated_adiabatic_step(prev.tmpc, prev.pres, lvl.pres)
        step = %{pres: lvl.pres, tmpc: t_next, env_t: lvl.tmpc, height: lvl.hght}
        {step, %{pres: lvl.pres, tmpc: t_next}}
      end)

    integral =
      traj
      |> Enum.chunk_every(2, 1, :discard)
      |> Enum.reduce(0.0, fn [a, b], acc ->
        diff = (a.env_t - a.tmpc + b.env_t - b.tmpc) / 2.0
        dlnp = :math.log(b.pres / a.pres)

        if diff > 0, do: acc + @r_d * diff * dlnp, else: acc
      end)

    round_j(integral)
  end

  defp theta_e(t_c, td_c, p_mb) do
    t_k = t_c + 273.15
    e = saturation_vapor_pressure_mb(td_c)
    r = @epsilon * e / max(p_mb - e, 1.0e-3)
    theta = t_k * :math.pow(1000.0 / p_mb, @r_d / @c_p)
    theta * :math.exp(@l_v * r / (@c_p * t_k))
  end

  # ── Lapse rates ─────────────────────────────────────────────────────

  defp lapse_rate_height(sorted_desc, lo_m, hi_m) do
    case {temp_at_height(sorted_desc, lo_m), temp_at_height(sorted_desc, hi_m)} do
      {t_lo, t_hi} when is_number(t_lo) and is_number(t_hi) ->
        Float.round((t_lo - t_hi) / ((hi_m - lo_m) / 1000.0), 2)

      _ ->
        nil
    end
  end

  defp lapse_rate_pressures(sorted_desc, p_lo, p_hi) do
    t_lo = pressure_lookup_temp(sorted_desc, p_lo)
    t_hi = pressure_lookup_temp(sorted_desc, p_hi)
    h_lo = height_at_pressure(sorted_desc, p_lo)
    h_hi = height_at_pressure(sorted_desc, p_hi)

    cond do
      not (is_number(t_lo) and is_number(t_hi)) -> nil
      not (is_number(h_lo) and is_number(h_hi)) -> nil
      h_hi == h_lo -> nil
      true -> Float.round((t_lo - t_hi) / ((h_hi - h_lo) / 1000.0), 2)
    end
  end

  defp temp_at_height(sorted_desc, target_m) do
    sorted_desc
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.find_value(fn [a, b] ->
      if a.hght <= target_m and target_m <= b.hght and b.hght != a.hght do
        f = (target_m - a.hght) / (b.hght - a.hght)
        a.tmpc + f * (b.tmpc - a.tmpc)
      end
    end)
  end

  defp height_at_pressure(_sorted, nil), do: nil

  defp height_at_pressure(sorted_desc, target_p) do
    sorted_desc
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.find_value(fn [a, b] ->
      if target_p <= a.pres and target_p >= b.pres and a.pres != b.pres do
        f = (a.pres - target_p) / (a.pres - b.pres)
        a.hght + f * (b.hght - a.hght)
      end
    end)
  end

  # ── Profile shape normaliser ────────────────────────────────────────

  defp normalize_level(level) when is_map(level) do
    pres = number(get_any(level, ["pres", "pres_mb", :pres, :pres_mb]))
    hght = number(get_any(level, ["hght", "hght_m", :hght, :hght_m]))
    tmpc = number(get_any(level, ["tmpc", :tmpc]))
    dwpc = number(get_any(level, ["dwpc", :dwpc]))

    if is_number(pres) and is_number(hght) and is_number(tmpc) do
      %{pres: pres, hght: hght, tmpc: tmpc, dwpc: dwpc}
    end
  end

  defp normalize_level(_), do: nil

  defp get_any(_map, []), do: nil

  defp get_any(map, [key | rest]) do
    case Map.fetch(map, key) do
      {:ok, v} -> v
      :error -> get_any(map, rest)
    end
  end

  defp number(nil), do: nil
  defp number(x) when is_integer(x), do: x * 1.0
  defp number(x) when is_float(x), do: x
  defp number(_), do: nil

  defp round_j(nil), do: nil
  defp round_j(v) when is_number(v), do: Float.round(v * 1.0, 1)
end