zen_quant

defmodule ZenQuant.Volatility do
  @moduledoc """
  Volatility calculations for trading systems.

  Pure functions for calculating historical (realized) volatility,
  comparing implied vs realized volatility, and computing volatility metrics.

  ## Example

      prices = [100, 102, 99, 103, 101, 105, 102]

      ZenQuant.Volatility.realized(prices)
      # => 0.023  (2.3% daily volatility)

      ZenQuant.Volatility.realized(prices, annualize: true)
      # => 0.365  (36.5% annualized)

  """

  use Descripex, namespace: "/volatility"

  @trading_days_per_year 365
  @yang_zhang_k_numerator 0.34
  @yang_zhang_k_denominator 1.34

  @typedoc "Volatility as decimal (e.g., 0.25 = 25%)"
  @type volatility :: float()

  @typedoc "IV comparison output format"
  @type iv_format :: :ratio | :premium | :premium_pct

  @typedoc "OHLC candle for volatility estimation"
  @type candle :: %{
          optional(:timestamp) => number(),
          optional(:open) => number(),
          :high => number(),
          :low => number(),
          optional(:close) => number(),
          optional(:volume) => number()
        }

  @typedoc "Positional OHLCV row returned by Bourse 0.6.x"
  @type ohlcv_row :: [number()]

  api(:normalize_candles, "Normalize positional OHLCV rows into candle maps.",
    params: [
      rows: [
        kind: :exchange_data,
        source: "fetch_ohlcv(symbol)",
        description: "List of [timestamp, open, high, low, close, volume] rows"
      ]
    ],
    returns: %{type: :list, description: "List of candle maps with atom keys"},
    returns_example: [
      %{timestamp: 1_700_000_000_000, open: 100.0, high: 105.0, low: 98.0, close: 103.0, volume: 42.0}
    ]
  )

  @spec normalize_candles([ohlcv_row()]) :: [candle()]
  def normalize_candles(rows) when is_list(rows) do
    Enum.map(rows, fn [timestamp, open, high, low, close, volume] ->
      %{timestamp: timestamp, open: open, high: high, low: low, close: close, volume: volume}
    end)
  end

  api(:realized, "Calculate realized (historical) volatility from close-to-close returns.",
    params: [
      prices: [kind: :value, description: "List of prices (chronological order, oldest first)"]
    ],
    opts: [
      annualize: [type: :boolean, default: false, description: "Whether to annualize the result"],
      trading_days: [type: :integer, default: 365, description: "Trading days per year for annualization"]
    ],
    returns: %{
      type: :float,
      description: "Standard deviation of returns, nil if insufficient data, or {:error, :invalid_prices}"
    },
    returns_example: 0.1095,
    errors: [:invalid_prices]
  )

  @spec realized([number()], keyword()) :: volatility() | nil | {:error, :invalid_prices}
  def realized(prices, opts \\ []) when is_list(prices) do
    cond do
      Enum.count_until(prices, 3) < 3 -> nil
      Enum.any?(prices, &(&1 <= 0)) -> {:error, :invalid_prices}
      true -> calculate_realized(prices, opts)
    end
  end

  # Computes standard deviation of log returns, optionally annualized
  defp calculate_realized(prices, opts) do
    returns = calculate_returns(prices)
    std_dev = standard_deviation(returns)

    if Keyword.get(opts, :annualize, false) do
      trading_days = Keyword.get(opts, :trading_days, @trading_days_per_year)
      std_dev * :math.sqrt(trading_days)
    else
      std_dev
    end
  end

  # Converts price series to log returns
  defp calculate_returns(prices) do
    prices
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.map(fn [prev, curr] -> :math.log(curr / prev) end)
  end

  api(:parkinson, "Calculate Parkinson volatility estimator from high-low ranges.",
    params: [
      candles: [
        kind: :exchange_data,
        source: "fetch_ohlcv(symbol)",
        description: "List of candles with :high and :low fields"
      ]
    ],
    opts: [
      annualize: [type: :boolean, default: false, description: "Whether to annualize"],
      trading_days: [type: :integer, default: 365, description: "Trading days per year"]
    ],
    returns: %{
      type: :float,
      description: "Parkinson volatility estimate, nil if insufficient data, or {:error, :invalid_prices}"
    },
    returns_example: 0.1095,
    errors: [:invalid_prices]
  )

  @spec parkinson([candle()], keyword()) :: volatility() | nil | {:error, :invalid_prices}
  def parkinson(candles, opts \\ []) when is_list(candles) do
    cond do
      Enum.count_until(candles, 2) < 2 -> nil
      has_invalid_high_low?(candles) -> {:error, :invalid_prices}
      true -> calculate_parkinson(candles, opts)
    end
  end

  # Checks if any candle has zero or negative high/low values
  defp has_invalid_high_low?(candles) do
    Enum.any?(candles, fn candle ->
      high = candle[:high] || candle["high"]
      low = candle[:low] || candle["low"]
      (is_number(high) and high <= 0) or (is_number(low) and low <= 0)
    end)
  end

  # Computes Parkinson estimator from high-low ranges
  defp calculate_parkinson(candles, opts) do
    n = length(candles)

    sum_squared =
      Enum.sum_by(candles, fn candle ->
        high = candle[:high] || candle["high"]
        low = candle[:low] || candle["low"]
        :math.pow(:math.log(high / low), 2)
      end)

    # Parkinson constant: 1 / (4 * ln(2))
    parkinson_constant = 1 / (4 * :math.log(2))
    variance = parkinson_constant * sum_squared / n
    std_dev = :math.sqrt(variance)

    if Keyword.get(opts, :annualize, false) do
      trading_days = Keyword.get(opts, :trading_days, @trading_days_per_year)
      std_dev * :math.sqrt(trading_days)
    else
      std_dev
    end
  end

  api(:iv_percentile, "Calculate IV percentile relative to historical distribution.",
    params: [
      current_iv: [kind: :value, description: "Current implied volatility"],
      historical_ivs: [kind: :value, description: "List of historical IV values"]
    ],
    returns: %{type: :float, description: "Percentile (0-100) of historical values below current, or nil if empty"},
    returns_example: 0.1095
  )

  @spec iv_percentile(number(), [number()]) :: float() | nil
  def iv_percentile(current_iv, historical_ivs) when is_number(current_iv) and is_list(historical_ivs) do
    if Enum.empty?(historical_ivs) do
      nil
    else
      count_below = Enum.count(historical_ivs, &(&1 < current_iv))
      count_below / length(historical_ivs) * 100
    end
  end

  api(:iv_rank, "Calculate IV rank (normalized position in min-max range).",
    params: [
      current_iv: [kind: :value, description: "Current implied volatility"],
      historical_ivs: [kind: :value, description: "List of historical IV values"]
    ],
    returns: %{type: :float, description: "Rank (0-100) showing position in min-max range, or nil if insufficient data"},
    returns_example: 0.1095
  )

  @spec iv_rank(number(), [number()]) :: float() | nil
  def iv_rank(current_iv, historical_ivs) when is_number(current_iv) and is_list(historical_ivs) do
    if Enum.count_until(historical_ivs, 2) < 2 do
      nil
    else
      {min_iv, max_iv} = Enum.min_max(historical_ivs)
      range = max_iv - min_iv

      if range == 0 do
        50.0
      else
        (current_iv - min_iv) / range * 100
      end
    end
  end

  api(:iv_vs_rv, "Compare implied volatility to realized volatility.",
    params: [
      implied_vol: [kind: :value, description: "Current implied volatility"],
      realized_vol: [kind: :value, description: "Realized volatility over same period"]
    ],
    opts: [
      format: [
        type: :atom,
        default: :ratio,
        description: ":ratio (IV/RV), :premium (IV-RV), or :premium_pct ((IV-RV)/RV*100)"
      ]
    ],
    returns: %{type: :float, description: "Comparison result in specified format, or nil if RV is zero"},
    returns_example: 0.1095
  )

  @spec iv_vs_rv(volatility(), volatility(), iv_format()) :: float() | nil
  def iv_vs_rv(implied_vol, realized_vol, format \\ :ratio) when is_number(implied_vol) and is_number(realized_vol) do
    if realized_vol == 0 do
      nil
    else
      calculate_iv_vs_rv(implied_vol, realized_vol, format)
    end
  end

  # Computes IV vs RV comparison in specified format (ratio, premium, or premium_pct)
  defp calculate_iv_vs_rv(implied_vol, realized_vol, :ratio) do
    implied_vol / realized_vol
  end

  defp calculate_iv_vs_rv(implied_vol, realized_vol, :premium) do
    implied_vol - realized_vol
  end

  defp calculate_iv_vs_rv(implied_vol, realized_vol, :premium_pct) do
    (implied_vol - realized_vol) / realized_vol * 100
  end

  api(:cone, "Calculate volatility cone for term structure analysis.",
    params: [
      prices: [kind: :value, description: "List of prices (chronological order)"],
      periods: [kind: :value, description: "List of lookback periods to calculate"]
    ],
    returns: %{type: :map, description: "Map of period => volatility value (or nil if insufficient data)"},
    returns_example: %{7 => 0.42, 30 => 0.36}
  )

  @spec cone([number()], [pos_integer()]) :: %{pos_integer() => volatility() | nil}
  def cone(prices, periods) when is_list(prices) and is_list(periods) do
    price_count = length(prices)

    periods
    |> Enum.filter(&(&1 <= price_count))
    |> Map.new(fn period ->
      recent_prices = Enum.take(prices, -period)
      {period, result_to_value(realized(recent_prices))}
    end)
  end

  api(:elevated?, "Check if volatility regime is elevated.",
    params: [
      current_vol: [kind: :value, description: "Current volatility"],
      baseline_vol: [kind: :value, description: "Normal/baseline volatility"]
    ],
    opts: [
      threshold: [type: :number, default: 1.5, description: "Multiple of baseline considered elevated"]
    ],
    returns: %{type: :boolean, description: "true if current_vol > baseline_vol * threshold"},
    returns_example: true
  )

  @spec elevated?(number(), number(), number()) :: boolean()
  def elevated?(current_vol, baseline_vol, threshold \\ 1.5)
      when is_number(current_vol) and is_number(baseline_vol) and is_number(threshold) do
    current_vol > baseline_vol * threshold
  end

  api(:garman_klass, "Calculate Garman-Klass volatility estimator from OHLC data.",
    params: [
      candles: [
        kind: :exchange_data,
        source: "fetch_ohlcv(symbol)",
        description: "List of candles with :open, :high, :low, :close fields"
      ]
    ],
    opts: [
      annualize: [type: :boolean, default: false, description: "Whether to annualize"],
      trading_days: [type: :integer, default: 365, description: "Trading days per year"]
    ],
    returns: %{
      type: :float,
      description: "Garman-Klass volatility, nil if insufficient data, or {:error, :invalid_prices}"
    },
    returns_example: 0.1095,
    errors: [:invalid_prices]
  )

  @spec garman_klass([candle()], keyword()) :: volatility() | nil | {:error, :invalid_prices}
  def garman_klass(candles, opts \\ []) when is_list(candles) do
    cond do
      Enum.count_until(candles, 2) < 2 -> nil
      has_invalid_ohlc?(candles) -> {:error, :invalid_prices}
      true -> calculate_garman_klass(candles, opts)
    end
  end

  # Checks if any candle has zero or negative OHLC values
  defp has_invalid_ohlc?(candles) do
    Enum.any?(candles, &candle_has_invalid_ohlc?/1)
  end

  # Checks a single candle for invalid OHLC values
  defp candle_has_invalid_ohlc?(candle) do
    Enum.any?(
      [
        candle[:open] || candle["open"],
        candle[:high] || candle["high"],
        candle[:low] || candle["low"],
        candle[:close] || candle["close"]
      ],
      &invalid_price?/1
    )
  end

  # Returns true if the value is a number <= 0
  defp invalid_price?(value) when is_number(value), do: value <= 0
  defp invalid_price?(_), do: false

  # Computes Garman-Klass estimator using OHLC data
  defp calculate_garman_klass(candles, opts) do
    n = length(candles)

    sum = Enum.sum_by(candles, &gk_single_candle/1)

    variance = sum / n
    std_dev = :math.sqrt(variance)

    if Keyword.get(opts, :annualize, false) do
      trading_days = Keyword.get(opts, :trading_days, @trading_days_per_year)
      std_dev * :math.sqrt(trading_days)
    else
      std_dev
    end
  end

  # Applies Garman-Klass formula to a single OHLC candle
  defp gk_single_candle(candle) do
    open = candle[:open] || candle["open"]
    high = candle[:high] || candle["high"]
    low = candle[:low] || candle["low"]
    close = candle[:close] || candle["close"]

    log_hl = :math.log(high / low)
    log_co = :math.log(close / open)

    # Garman-Klass formula
    0.5 * :math.pow(log_hl, 2) - (2 * :math.log(2) - 1) * :math.pow(log_co, 2)
  end

  api(:yang_zhang, "Calculate Yang-Zhang volatility estimator from OHLC data.",
    params: [
      candles: [
        kind: :exchange_data,
        source: "fetch_ohlcv(symbol)",
        description: "List of consecutive candles with :open, :high, :low, :close fields"
      ]
    ],
    opts: [
      annualize: [type: :boolean, default: false, description: "Whether to annualize"],
      trading_days: [type: :integer, default: 365, description: "Trading days per year"]
    ],
    returns: %{
      type: :float,
      description: "Yang-Zhang volatility, nil if insufficient data, or {:error, :invalid_prices}"
    },
    returns_example: 0.1095,
    errors: [:invalid_prices]
  )

  @spec yang_zhang([candle()], keyword()) :: volatility() | nil | {:error, :invalid_prices}
  def yang_zhang(candles, opts \\ []) when is_list(candles) do
    cond do
      Enum.count_until(candles, 3) < 3 -> nil
      has_invalid_ohlc?(candles) -> {:error, :invalid_prices}
      true -> calculate_yang_zhang(candles, opts)
    end
  end

  # Computes Yang-Zhang estimator by combining overnight, open-close, and Rogers-Satchell terms
  defp calculate_yang_zhang(candles, opts) do
    n = length(candles)
    overnight_returns = yz_overnight_returns(candles)
    open_close_returns = yz_open_close_returns(candles)
    overnight_variance = mean_centered_variance(overnight_returns, n)
    open_close_variance = mean_centered_variance(open_close_returns, n)
    rs_variance = yz_rogers_satchell_variance(candles, n)
    k = @yang_zhang_k_numerator / (@yang_zhang_k_denominator + (n + 1) / (n - 1))
    variance = overnight_variance + k * open_close_variance + (1 - k) * rs_variance
    std_dev = :math.sqrt(variance)

    if Keyword.get(opts, :annualize, false) do
      trading_days = Keyword.get(opts, :trading_days, @trading_days_per_year)
      std_dev * :math.sqrt(trading_days)
    else
      std_dev
    end
  end

  # Computes close-to-open overnight returns from consecutive candles
  defp yz_overnight_returns(candles) do
    candles
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.map(fn [previous, current] ->
      prev_close = previous[:close] || previous["close"]
      curr_open = current[:open] || current["open"]
      :math.log(curr_open / prev_close)
    end)
  end

  # Computes open-to-close returns for each candle
  defp yz_open_close_returns(candles) do
    Enum.map(candles, fn candle ->
      open = candle[:open] || candle["open"]
      close = candle[:close] || candle["close"]
      :math.log(close / open)
    end)
  end

  # Computes Rogers-Satchell component variance
  defp yz_rogers_satchell_variance(candles, n) do
    Enum.sum_by(candles, &yz_rogers_satchell_single/1) / n
  end

  # Computes Rogers-Satchell term for a single candle
  defp yz_rogers_satchell_single(candle) do
    open = candle[:open] || candle["open"]
    high = candle[:high] || candle["high"]
    low = candle[:low] || candle["low"]
    close = candle[:close] || candle["close"]

    :math.log(high / close) * :math.log(high / open) + :math.log(low / close) * :math.log(low / open)
  end

  # Computes mean-centered variance using provided divisor
  defp mean_centered_variance([], _divisor), do: 0.0

  defp mean_centered_variance(values, divisor) do
    mean = Enum.sum(values) / length(values)

    Enum.sum_by(values, &:math.pow(&1 - mean, 2)) / divisor
  end

  api(:rolling, "Calculate rolling volatility over a price series.",
    params: [
      prices: [kind: :value, description: "List of prices (chronological order)"],
      window: [kind: :value, description: "Rolling window size (> 2)"]
    ],
    opts: [
      annualize: [type: :boolean, default: false, description: "Whether to annualize"],
      trading_days: [type: :integer, default: 365, description: "Trading days per year"]
    ],
    returns: %{type: :list, description: "List of volatility values, empty if insufficient data"},
    returns_example: [0.38, 0.41]
  )

  @spec rolling([number()], pos_integer(), keyword()) :: [volatility()]
  def rolling(prices, window, opts \\ []) when is_list(prices) and is_integer(window) and window > 2 do
    if length(prices) < window do
      []
    else
      calculate_rolling(prices, window, opts)
    end
  end

  # Calculates rolling volatility using sliding window
  defp calculate_rolling(prices, window, opts) do
    prices
    |> Enum.chunk_every(window, 1, :discard)
    |> Enum.map(&realized(&1, opts))
    |> Enum.map(&result_to_value/1)
    |> Enum.reject(&is_nil/1)
  end

  # Converts error tuples to nil for aggregate functions
  defp result_to_value({:error, _}), do: nil
  defp result_to_value(value), do: value

  # Calculates sample standard deviation of a series (n-1 for unbiased estimator)
  defp standard_deviation([]), do: 0.0
  defp standard_deviation([_]), do: 0.0

  defp standard_deviation(values) do
    n = length(values)
    mean = Enum.sum(values) / n

    variance = Enum.sum_by(values, &:math.pow(&1 - mean, 2)) / (n - 1)

    :math.sqrt(variance)
  end
end