ex_fairness

defmodule ExFairness.Report do
  @moduledoc """
  Fairness report generation and export.

  Provides comprehensive fairness assessment across multiple metrics with
  multiple export formats (Markdown, JSON).

  ## Examples

      iex> predictions = Nx.tensor([1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1])
      iex> labels = Nx.tensor([1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
      iex> sensitive = Nx.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
      iex> report = ExFairness.Report.generate(predictions, labels, sensitive)
      iex> Map.has_key?(report, :overall_assessment)
      true

  """

  alias ExFairness.Metrics.{
    Calibration,
    DemographicParity,
    EqualizedOdds,
    EqualOpportunity,
    PredictiveParity
  }

  @all_metrics [
    :demographic_parity,
    :equalized_odds,
    :equal_opportunity,
    :predictive_parity,
    :calibration
  ]

  @type report :: %{
          optional(:demographic_parity) => DemographicParity.result(),
          optional(:equalized_odds) => EqualizedOdds.result(),
          optional(:equal_opportunity) => EqualOpportunity.result(),
          optional(:predictive_parity) => PredictiveParity.result(),
          optional(:calibration) => Calibration.result(),
          overall_assessment: String.t(),
          passed_count: non_neg_integer(),
          failed_count: non_neg_integer(),
          total_count: non_neg_integer()
        }

  @doc """
  Generates a comprehensive fairness report across multiple metrics.

  ## Parameters

    * `predictions` - Binary predictions tensor (0 or 1)
    * `labels` - Binary labels tensor (0 or 1)
    * `sensitive_attr` - Binary sensitive attribute tensor (0 or 1)
    * `opts` - Options:
      * `:metrics` - List of metrics to include (default: all available, calibration only when `:probabilities` is provided)
      * `:threshold` - Fairness threshold to pass to all metrics
      * `:min_per_group` - Minimum samples per group
      * `:probabilities` - Predicted probabilities (required for `:calibration`)
      * CI/testing options forwarded to metrics:
        * `:include_ci` - Enable bootstrap confidence intervals
        * `:bootstrap_samples`, `:confidence_level`, `:stratified`
        * `:statistical_test` - e.g., :z_test | :chi_square | :permutation
        * `:alpha`, `:n_permutations`

  ## Returns

  A map containing:
    * Metric results (one key per requested metric)
    * `:overall_assessment` - Summary of fairness across all metrics
    * `:passed_count` - Number of metrics that passed
    * `:failed_count` - Number of metrics that failed
    * `:total_count` - Total number of metrics evaluated

  ## Examples

      iex> predictions = Nx.tensor([1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1])
      iex> labels = Nx.tensor([1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
      iex> sensitive = Nx.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
      iex> report = ExFairness.Report.generate(predictions, labels, sensitive, metrics: [:demographic_parity])
      iex> report.total_count
      1

  """
  @spec generate(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), keyword()) :: report()
  def generate(predictions, labels, sensitive_attr, opts \\ []) do
    metrics = Keyword.get(opts, :metrics, default_metrics(opts))

    # Compute each requested metric
    results =
      Enum.reduce(metrics, %{}, fn metric, acc ->
        result = compute_metric(metric, predictions, labels, sensitive_attr, opts)
        Map.put(acc, metric, result)
      end)

    # Count passes/failures
    passed_count = Enum.count(results, fn {_metric, result} -> result.passes end)
    failed_count = Enum.count(results, fn {_metric, result} -> !result.passes end)
    total_count = map_size(results)

    # Generate overall assessment
    overall_assessment = generate_overall_assessment(passed_count, failed_count, total_count)

    results
    |> Map.put(:overall_assessment, overall_assessment)
    |> Map.put(:passed_count, passed_count)
    |> Map.put(:failed_count, failed_count)
    |> Map.put(:total_count, total_count)
  end

  @doc """
  Exports a fairness report to Markdown format.

  ## Parameters

    * `report` - A fairness report generated by `generate/4`

  ## Returns

  A Markdown-formatted string containing the report.

  ## Examples

      iex> predictions = Nx.tensor([1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1])
      iex> labels = Nx.tensor([1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
      iex> sensitive = Nx.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
      iex> report = ExFairness.Report.generate(predictions, labels, sensitive, metrics: [:demographic_parity])
      iex> markdown = ExFairness.Report.to_markdown(report)
      iex> String.contains?(markdown, "# Fairness Report")
      true

  """
  @spec to_markdown(report()) :: String.t()
  def to_markdown(report) do
    """
    # Fairness Report

    ## Overall Assessment

    #{report.overall_assessment}

    **Summary:** #{report.passed_count} of #{report.total_count} metrics passed.

    ## Metric Results

    | Metric | Passes | Disparity | Threshold |
    |--------|--------|-----------|-----------|
    #{format_metrics_table(report)}

    ## Detailed Results

    #{format_detailed_results(report)}
    """
  end

  @doc """
  Exports a fairness report to JSON format.

  ## Parameters

    * `report` - A fairness report generated by `generate/4`

  ## Returns

  A JSON-formatted string containing the report.

  ## Examples

      iex> predictions = Nx.tensor([1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1])
      iex> labels = Nx.tensor([1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
      iex> sensitive = Nx.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
      iex> report = ExFairness.Report.generate(predictions, labels, sensitive, metrics: [:demographic_parity])
      iex> json = ExFairness.Report.to_json(report)
      iex> String.starts_with?(json, "{")
      true

  """
  @spec to_json(report()) :: String.t()
  def to_json(report) do
    Jason.encode!(report, pretty: true)
  end

  # Private functions

  @spec compute_metric(atom(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), keyword()) :: map()
  defp compute_metric(:demographic_parity, predictions, _labels, sensitive_attr, opts) do
    DemographicParity.compute(predictions, sensitive_attr, opts)
  end

  defp compute_metric(:equalized_odds, predictions, labels, sensitive_attr, opts) do
    EqualizedOdds.compute(predictions, labels, sensitive_attr, opts)
  end

  defp compute_metric(:equal_opportunity, predictions, labels, sensitive_attr, opts) do
    EqualOpportunity.compute(predictions, labels, sensitive_attr, opts)
  end

  defp compute_metric(:predictive_parity, predictions, labels, sensitive_attr, opts) do
    PredictiveParity.compute(predictions, labels, sensitive_attr, opts)
  end

  defp compute_metric(:calibration, _predictions, labels, sensitive_attr, opts) do
    probabilities =
      Keyword.get(opts, :probabilities) ||
        raise ArgumentError,
              ":probabilities option is required when requesting :calibration in fairness_report/4"

    calibration = Calibration.compute(probabilities, labels, sensitive_attr, opts)

    reliability =
      if Keyword.get(opts, :include_reliability, true) do
        Calibration.reliability_diagram(probabilities, labels, sensitive_attr, opts)
      else
        nil
      end

    Map.put(calibration, :reliability_diagram, reliability)
  end

  @spec generate_overall_assessment(non_neg_integer(), non_neg_integer(), non_neg_integer()) ::
          String.t()
  defp generate_overall_assessment(passed, _failed, total) when passed == total do
    "✓ All #{total} fairness metrics passed. The model demonstrates fairness across all evaluated criteria."
  end

  defp generate_overall_assessment(0, failed, total) when failed == total do
    "✗ All #{total} fairness metrics failed. The model exhibits significant fairness concerns that require attention."
  end

  defp generate_overall_assessment(passed, failed, total) do
    "⚠ Mixed results: #{passed} of #{total} metrics passed, #{failed} failed. The model demonstrates fairness in some areas but has concerns in others."
  end

  @spec format_metrics_table(report()) :: String.t()
  defp format_metrics_table(report) do
    metrics = metrics_in_report(report)
    include_ci? = has_confidence_intervals?(report, metrics)
    include_p? = has_p_values?(report, metrics)

    header =
      ["Metric", "Passes", "Disparity", "Threshold"]
      |> maybe_append(include_ci?, "CI")
      |> maybe_append(include_p?, "p-value")
      |> Enum.join(" | ")

    separator =
      ["--------", "------", "---------", "----------"]
      |> maybe_append(include_ci?, "----")
      |> maybe_append(include_p?, "-------")
      |> Enum.join(" | ")

    rows =
      metrics
      |> Enum.map(fn metric ->
        result = Map.get(report, metric)
        passes = if result.passes, do: "✓", else: "✗"
        disparity = get_primary_disparity(metric, result)
        threshold = result.threshold

        base = [
          format_metric_name(metric),
          passes,
          Float.round(disparity, 3),
          Float.round(threshold, 3)
        ]

        base
        |> maybe_append(include_ci?, format_ci(result))
        |> maybe_append(include_p?, format_p_value(result))
        |> Enum.join(" | ")
      end)
      |> Enum.map(&("| " <> &1 <> " |"))
      |> Enum.join("\n")

    """
    | #{header} |
    | #{separator} |
    #{rows}
    """
  end

  @spec format_detailed_results(report()) :: String.t()
  defp format_detailed_results(report) do
    metrics_in_report(report)
    |> Enum.map(fn metric ->
      result = Map.get(report, metric)

      """
      ### #{format_metric_name(metric)}

      **Status:** #{if result.passes, do: "✓ Passed", else: "✗ Failed"}

      #{result.interpretation}
      #{format_ci_block(result)}
      #{format_p_value_block(result)}
      #{format_reliability(metric, result)}
      """
    end)
    |> Enum.join("\n")
  end

  @spec format_metric_name(atom()) :: String.t()
  defp format_metric_name(:demographic_parity), do: "Demographic Parity"
  defp format_metric_name(:equalized_odds), do: "Equalized Odds"
  defp format_metric_name(:equal_opportunity), do: "Equal Opportunity"
  defp format_metric_name(:predictive_parity), do: "Predictive Parity"
  defp format_metric_name(:calibration), do: "Calibration Fairness"

  @spec get_primary_disparity(atom(), map()) :: float()
  defp get_primary_disparity(:demographic_parity, result), do: result.disparity
  defp get_primary_disparity(:equal_opportunity, result), do: result.disparity
  defp get_primary_disparity(:predictive_parity, result), do: result.disparity

  defp get_primary_disparity(:equalized_odds, result) do
    max(result.tpr_disparity, result.fpr_disparity)
  end

  defp get_primary_disparity(:calibration, result), do: result.disparity

  defp metrics_in_report(report) do
    report
    |> Map.keys()
    |> Enum.filter(&(&1 in @all_metrics))
  end

  defp default_metrics(opts) do
    if Keyword.has_key?(opts, :probabilities) do
      @all_metrics
    else
      @all_metrics -- [:calibration]
    end
  end

  defp format_reliability(:calibration, %{reliability_diagram: diagram}) do
    if diagram do
      rows =
        diagram.bins
        |> Enum.map(fn bin ->
          {low, high} = bin.range

          [
            "#{bin.bin_index}",
            "[#{Float.round(low, 3)}, #{Float.round(high, 3)})",
            bin.group_a.count,
            Float.round(bin.group_a.confidence, 3),
            Float.round(bin.group_a.accuracy, 3),
            bin.group_b.count,
            Float.round(bin.group_b.confidence, 3),
            Float.round(bin.group_b.accuracy, 3)
          ]
          |> Enum.join(" | ")
        end)
        |> Enum.join("\n")

      """
      \n**Calibration Reliability Diagram** (#{diagram.strategy} bins = #{diagram.n_bins})

      | Bin | Range | A Count | A Conf | A Acc | B Count | B Conf | B Acc |
      |-----|-------|---------|--------|-------|---------|--------|-------|
      #{rows}
      """
    else
      ""
    end
  end

  defp format_reliability(_metric, _result), do: ""

  defp format_ci(result) do
    case Map.get(result, :confidence_interval) do
      {low, high} -> "[#{Float.round(low, 3)}, #{Float.round(high, 3)}]"
      _ -> "-"
    end
  end

  defp format_ci_block(result) do
    case Map.get(result, :confidence_interval) do
      {low, high} ->
        "\n**Confidence interval:** [#{Float.round(low, 4)}, #{Float.round(high, 4)}]"

      _ ->
        ""
    end
  end

  defp format_p_value(result) do
    case Map.get(result, :p_value) do
      p when is_number(p) -> Float.round(p, 4)
      _ -> "-"
    end
  end

  defp format_p_value_block(result) do
    case Map.get(result, :p_value) do
      p when is_number(p) ->
        sig = Map.get(result, :significant)
        conclusion = if is_boolean(sig), do: " (significant? #{sig})", else: ""
        "\n**p-value:** #{Float.round(p, 4)}#{conclusion}"

      _ ->
        ""
    end
  end

  defp has_confidence_intervals?(report, metrics) do
    Enum.any?(metrics, fn metric ->
      result = Map.get(report, metric)
      match?({_, _}, Map.get(result, :confidence_interval))
    end)
  end

  defp has_p_values?(report, metrics) do
    Enum.any?(metrics, fn metric ->
      result = Map.get(report, metric)
      is_number(Map.get(result, :p_value))
    end)
  end

  defp maybe_append(list, true, value), do: list ++ [value]
  defp maybe_append(list, false, _value), do: list
end