defmodule Microwaveprop.Propagation.Model do
  @moduledoc """
  Neural network model for microwave propagation prediction.

  Uses Axon to define a model that predicts propagation conditions
  (score 0-100 per band) from atmospheric and temporal features.

  ## Features (inputs)

  Atmospheric (from HRRR):
    - surface_temp_c — surface temperature
    - surface_dewpoint_c — surface dewpoint
    - surface_pressure_mb — surface pressure
    - abs_humidity — absolute humidity (g/m³), derived
    - td_depression — temp minus dewpoint (°C)
    - min_refractivity_gradient — minimum dN/dh from profile
    - hpbl_m — planetary boundary layer height
    - pwat_mm — precipitable water

  Temporal:
    - utc_hour — hour of day (0-23), sin/cos encoded
    - month — month of year (1-12), sin/cos encoded

  Band:
    - freq_mhz — operating frequency, log-scaled

  ## Target (output)

  - score — propagation score (0-100), normalized to 0-1 for training

  ## Architecture

  Feed-forward neural network with:
    - Input: 13 features (8 atmospheric + 4 cyclical temporal + 1 frequency)
    - Hidden: 2 dense layers with ReLU activation
    - Output: 1 sigmoid unit (score 0-1, scaled to 0-100)

  The model is intentionally simple to start. We can add complexity
  (LSTM for temporal sequences, attention, ensemble) once we have
  baseline performance metrics.
  """

  @feature_count 13
  @models_dir Path.join(:code.priv_dir(:microwaveprop), "models")
  @default_path Path.join(@models_dir, "propagation_v1.nx")

  @doc """
  Builds the Axon model graph. Does not initialize parameters.

  Returns an `%Axon{}` struct ready for `Axon.build/2` or `Axon.Loop.trainer/3`.
  """
  def build do
    "features"
    |> Axon.input(shape: {nil, @feature_count})
    |> Axon.dense(64, activation: :relu, name: "hidden_1")
    |> Axon.dropout(rate: 0.2, name: "dropout_1")
    |> Axon.dense(32, activation: :relu, name: "hidden_2")
    |> Axon.dropout(rate: 0.1, name: "dropout_2")
    |> Axon.dense(1, activation: :sigmoid, name: "output")
  end

  @doc """
  Initializes random model parameters.

  Returns `%{...}` map of layer name → tensor params.
  """
  def init do
    model = build()
    {init_fn, _predict_fn} = Axon.build(model)
    template = Nx.template({1, @feature_count}, :f32)
    init_fn.(template, Axon.ModelState.empty())
  end

  @doc """
  Saves model parameters to disk. Defaults to `priv/models/propagation_v1.nx`.
  """
  def save(params, path \\ @default_path) do
    File.mkdir_p!(Path.dirname(path))
    binary = Nx.serialize(params)
    File.write!(path, binary)
    :ok
  end

  @doc """
  Loads model parameters from disk. Returns `{:ok, params}` or `:error`.
  """
  def load(path \\ @default_path) do
    case File.read(path) do
      {:ok, binary} -> {:ok, Nx.deserialize(binary)}
      {:error, _} -> :error
    end
  end

  @doc """
  Loads saved parameters if available, otherwise initializes random ones.
  """
  def load_or_init(path \\ @default_path) do
    case load(path) do
      {:ok, params} -> params
      :error -> init()
    end
  end

  @doc """
  Runs a forward pass with the given parameters and input features.

  `features` should be an `{batch_size, 13}` tensor of float32.
  Returns an `{batch_size, 1}` tensor of scores in [0, 1].
  """
  def predict(params, features) do
    model = build()
    {_init_fn, predict_fn} = Axon.build(model)
    predict_fn.(params, features)
  end

  @doc """
  Returns the ordered list of feature names expected by the model.
  """
  def feature_names do
    [
      :surface_temp_c,
      :surface_dewpoint_c,
      :surface_pressure_mb,
      :abs_humidity,
      :td_depression,
      :min_refractivity_gradient,
      :hpbl_m,
      :pwat_mm,
      :utc_hour_sin,
      :utc_hour_cos,
      :month_sin,
      :month_cos,
      :log_freq_mhz
    ]
  end

  @doc """
  Encodes raw condition data into a feature tensor row.

  Takes a map of raw values and returns a flat list of 13 floats
  ready to be stacked into a batch tensor.
  """
  def encode_features(%{} = conditions) do
    temp_c = conditions[:surface_temp_c] || 20.0
    dewpoint_c = conditions[:surface_dewpoint_c] || 10.0
    pressure_mb = conditions[:surface_pressure_mb] || 1013.0
    grad = conditions[:min_refractivity_gradient] || -70.0
    hpbl = conditions[:hpbl_m] || 500.0
    pwat = conditions[:pwat_mm] || 20.0
    utc_hour = conditions[:utc_hour] || 12
    month = conditions[:month] || 6
    freq_mhz = conditions[:freq_mhz] || 10_000

    abs_humidity = abs_humidity(temp_c, dewpoint_c)
    td_depression = temp_c - dewpoint_c

    # Cyclical encoding for time features
    hour_rad = 2 * :math.pi() * utc_hour / 24
    month_rad = 2 * :math.pi() * (month - 1) / 12

    [
      temp_c,
      dewpoint_c,
      pressure_mb,
      abs_humidity,
      td_depression,
      grad,
      hpbl,
      pwat,
      :math.sin(hour_rad),
      :math.cos(hour_rad),
      :math.sin(month_rad),
      :math.cos(month_rad),
      :math.log(freq_mhz)
    ]
  end

  defp abs_humidity(temp_c, dewpoint_c) do
    e_sat = 6.112 * :math.exp(17.67 * dewpoint_c / (dewpoint_c + 243.5))
    217.0 * e_sat / (temp_c + 273.15)
  end
end