anova_visualizer

defmodule AnovaVisualizer do
  @moduledoc """
  Generates interactive VegaLite visualizations for ANOVA and Tukey's HSD analysis results
  using Kino.VegaLite for Livebook integration.
  """

  @doc """
  Creates a comprehensive visualization layout for ANOVA and Tukey's HSD results.

  ## Example

      # Run ANOVA and post-hoc test, then visualize:
      # groups
      # |> ANOVA.one_way()
      # |> TukeyHSD.test(0.05)
      # |> AnovaVisualizer.visualize()
  """
  def visualize(
        %{
          anova: anova_result,
          post_hoc_test: post_hoc_result
        },
        title \\ "ANOVA One-Way with Tukey's HSD Post-Hoc Test",
        groups \\ nil
      ) do
    Kino.Layout.grid(
      [
        render_header(title),
        render_anova_summary(anova_result),
        render_anova_table(anova_result),
        render_group_means_chart(anova_result, groups),
        render_posthoc_summary(post_hoc_result, groups),
        render_pairwise_comparisons(post_hoc_result, groups),
        render_confidence_intervals(post_hoc_result, groups),
        render_effect_size_comparison(post_hoc_result, groups)
      ],
      columns: 1
    )
  end

  defp render_header(title) do
    content = """
    # #{title}

    ---
    """

    Kino.Markdown.new(content)
  end

  defp render_anova_summary(anova) do
    significance =
      if anova.test_results.p_value < 0.05, do: "✅ **Significant**", else: "❌ Not Significant"

    """
    ## 🔬 ANOVA Summary Statistics

    | Metric | Value |
    |--------|-------|
    | **Number of Groups** | #{anova.summary.groups} |
    | **Total Observations** | #{anova.summary.total_observations} |
    | **Overall Mean** | #{format_number(anova.summary.overall_mean)} |
    | **F-Statistic** | #{format_number(anova.test_results.f_statistic)} |
    | **p-Value** | #{format_scientific(anova.test_results.p_value)} |
    | **Significance** | #{significance} |
    | **Eta Squared (η²)** | #{format_number(anova.test_results.eta_squared)} (#{interpret_eta_squared(anova.test_results.eta_squared)}) |
    | **Omega Squared (ω²)** | #{format_number(anova.test_results.omega_squared)} |

    ---
    """
    |> Kino.Markdown.new()
  end

  defp render_anova_table(anova) do
    """
    ## ANOVA Table

    | Source | Sum of Squares | df | Mean Square | F-Statistic | p-Value |
    |--------|----------------|-------|-------------|-------------|---------|
    | **Between Groups** | #{format_number(anova.anova_table.between.ss)} | #{anova.anova_table.between.df} | #{format_number(anova.anova_table.between.ms)} | #{format_number(anova.test_results.f_statistic)} | #{format_scientific(anova.test_results.p_value)} |
    | **Within Groups** | #{format_number(anova.anova_table.within.ss)} | #{anova.anova_table.within.df} | #{format_number(anova.anova_table.within.ms)} | — | — |
    | **Total** | #{format_number(anova.anova_table.total.ss)} | #{anova.anova_table.total.df} | — | — | — |

    ---
    """
    |> Kino.Markdown.new()
  end

  defp group_name(nil, idx), do: "Group #{idx}"
  defp group_name(groups, idx), do: Enum.at(groups, idx - 1) || "Group #{idx}"

  defp render_group_means_chart(anova, groups) do
    data =
      anova.summary.group_means
      |> Enum.with_index(1)
      |> Enum.map(fn {mean, idx} ->
        %{
          "group" => group_name(groups, idx),
          "mean" => mean,
          "group_num" => idx
        }
      end)

    VegaLite.new(width: 600, height: 400, title: "Group Means Comparison")
    |> VegaLite.data_from_values(data)
    |> VegaLite.mark(:bar, tooltip: true)
    |> VegaLite.encode_field(:x, "group",
      type: :nominal,
      title: "Group",
      axis: [label_angle: 0, label_font_size: 14]
    )
    |> VegaLite.encode_field(:y, "mean",
      type: :quantitative,
      title: "Mean Value",
      scale: [zero: true]
    )
    |> VegaLite.encode_field(:color, "group_num",
      type: :nominal,
      scale: [range: ["#667eea", "#764ba2"]],
      legend: false
    )
    |> VegaLite.encode(:tooltip, [
      [field: "group", type: :nominal, title: "Group"],
      [field: "mean", type: :quantitative, title: "Mean", format: ".2f"]
    ])
    |> Kino.VegaLite.new()
  end

  defp render_posthoc_summary(post_hoc, groups) do
    significant_pairs =
      post_hoc.summary.significant_pairs
      |> Enum.map(fn {g1, g2} -> "#{group_name(groups, g1)} vs #{group_name(groups, g2)}" end)
      |> Enum.join(", ")

    """
    ## Tukey's HSD Post-Hoc Analysis Summary

    | Metric | Value |
    |--------|-------|
    | **Total Comparisons** | #{post_hoc.summary.total_comparisons} |
    | **Significant Comparisons** | #{post_hoc.summary.significant_comparisons} |
    | **Non-Significant Comparisons** | #{post_hoc.summary.non_significant_comparisons} |
    | **Significant Pairs** | #{significant_pairs} |
    | **Mean Difference (avg)** | #{format_number(post_hoc.summary.difference_stats.mean)} |
    | **Mean Effect Size** | #{format_number(post_hoc.summary.effect_size_stats.mean)} (#{interpret_cohens_d(post_hoc.summary.effect_size_stats.mean)}) |

    ---
    """
    |> Kino.Markdown.new()
  end

  defp render_pairwise_comparisons(post_hoc, groups) do
    comparisons = post_hoc.pairwise_comparisons

    headers =
      comparisons
      |> Enum.map(fn comp ->
        {g1, g2} = comp.groups
        "#{group_name(groups, g1)} vs #{group_name(groups, g2)}"
      end)

    header_row = "| **Statistic** | " <> Enum.join(headers, " | ") <> " |"
    separator = "| --- |" <> String.duplicate(" --- |", length(comparisons))

    ci_level = comparisons |> List.first() |> then(& &1.confidence_interval.level) |> trunc()

    rows =
      [
        {"**Significance**",
         Enum.map(comparisons, fn c ->
           if c.significant?, do: "✅", else: "❌"
         end)},
        {"**Mean Difference**", Enum.map(comparisons, &format_number(&1.difference))},
        {"**Group Means (A / B)**",
         Enum.map(comparisons, fn c ->
           {m1, m2} = c.means
           "#{format_number(m1)} / #{format_number(m2)}"
         end)},
        {"**Q-Statistic**", Enum.map(comparisons, &format_number(&1.q_statistic))},
        {"**p-Value**", Enum.map(comparisons, &format_scientific(&1.p_value))},
        {"**Standard Error**", Enum.map(comparisons, &format_number(&1.standard_error))},
        {"**Effect Size (Cohen's d)**",
         Enum.map(comparisons, fn c ->
           "#{format_number(c.effect_size)} (#{interpret_cohens_d(c.effect_size)})"
         end)},
        {"**#{ci_level}% CI Lower**",
         Enum.map(comparisons, &format_number(&1.confidence_interval.lower))},
        {"**#{ci_level}% CI Upper**",
         Enum.map(comparisons, &format_number(&1.confidence_interval.upper))}
      ]
      |> Enum.map(fn {label, values} ->
        "| #{label} | " <> Enum.join(values, " | ") <> " |"
      end)

    table = Enum.join([header_row, separator | rows], "\n")

    """
    ## Pairwise Comparisons

    #{table}

    ---
    """
    |> Kino.Markdown.new()
  end

  defp render_confidence_intervals(post_hoc, groups) do
    # Create data for bars (mean difference)
    bar_data =
      post_hoc.pairwise_comparisons
      |> Enum.map(fn comp ->
        {g1, g2} = comp.groups

        %{
          "comparison" => "#{group_name(groups, g1)} vs #{group_name(groups, g2)}",
          "difference" => comp.difference,
          "ci_lower" => comp.confidence_interval.lower,
          "ci_upper" => comp.confidence_interval.upper,
          "significant" => if(comp.significant?, do: "Significant", else: "Not Significant")
        }
      end)

    VegaLite.new(width: 700, height: 400, title: "Mean Differences with 95% Confidence Intervals")
    |> VegaLite.data_from_values(bar_data)
    |> VegaLite.layers([
      # Bar layer for mean difference
      VegaLite.new()
      |> VegaLite.mark(:bar, tooltip: true)
      |> VegaLite.encode_field(:x, "comparison",
        type: :nominal,
        title: "Comparison",
        axis: [label_angle: 0]
      )
      |> VegaLite.encode_field(:y, "difference",
        type: :quantitative,
        title: "Mean Difference",
        scale: [zero: true]
      )
      |> VegaLite.encode_field(:color, "significant",
        type: :nominal,
        scale: [
          domain: ["Significant", "Not Significant"],
          range: ["#10b981", "#ef4444"]
        ],
        legend: [title: "Significance"]
      ),

      # Error bar layer for confidence intervals
      VegaLite.new()
      |> VegaLite.mark(:rule, color: "#1f2937", size: 2)
      |> VegaLite.encode_field(:x, "comparison", type: :nominal)
      |> VegaLite.encode_field(:y, "ci_lower", type: :quantitative)
      |> VegaLite.encode_field(:y2, "ci_upper"),

      # Tick marks at CI bounds
      VegaLite.new()
      |> VegaLite.mark(:tick, color: "#1f2937", size: 15, thickness: 2)
      |> VegaLite.encode_field(:x, "comparison", type: :nominal)
      |> VegaLite.encode_field(:y, "ci_lower", type: :quantitative),
      VegaLite.new()
      |> VegaLite.mark(:tick, color: "#1f2937", size: 15, thickness: 2)
      |> VegaLite.encode_field(:x, "comparison", type: :nominal)
      |> VegaLite.encode_field(:y, "ci_upper", type: :quantitative)
    ])
    |> VegaLite.encode(:tooltip, [
      [field: "comparison", type: :nominal, title: "Comparison"],
      [field: "difference", type: :quantitative, title: "Mean Difference", format: ".2f"],
      [field: "ci_lower", type: :quantitative, title: "95% CI Lower", format: ".2f"],
      [field: "ci_upper", type: :quantitative, title: "95% CI Upper", format: ".2f"],
      [field: "significant", type: :nominal, title: "Significant?"]
    ])
    |> Kino.VegaLite.new()
  end

  defp render_effect_size_comparison(post_hoc, groups) do
    data =
      post_hoc.pairwise_comparisons
      |> Enum.map(fn comp ->
        {g1, g2} = comp.groups

        %{
          "comparison" => "#{group_name(groups, g1)} vs #{group_name(groups, g2)}",
          "effect_size" => comp.effect_size,
          "significant" => if(comp.significant?, do: "Significant", else: "Not Significant"),
          "interpretation" => interpret_cohens_d(comp.effect_size)
        }
      end)

    VegaLite.new(width: 600, height: 400, title: "Effect Sizes (Cohen's d) by Comparison")
    |> VegaLite.data_from_values(data)
    |> VegaLite.mark(:bar, tooltip: true)
    |> VegaLite.encode_field(:x, "comparison",
      type: :nominal,
      title: "Comparison",
      axis: [label_angle: -45]
    )
    |> VegaLite.encode_field(:y, "effect_size",
      type: :quantitative,
      title: "Effect Size (Cohen's d)",
      scale: [zero: true]
    )
    |> VegaLite.encode_field(:color, "significant",
      type: :nominal,
      scale: [
        domain: ["Significant", "Not Significant"],
        range: ["#10b981", "#ef4444"]
      ],
      legend: [title: "Significance"]
    )
    |> VegaLite.encode(:tooltip, [
      [field: "comparison", type: :nominal, title: "Comparison"],
      [field: "effect_size", type: :quantitative, title: "Effect Size", format: ".3f"],
      [field: "interpretation", type: :nominal, title: "Interpretation"],
      [field: "significant", type: :nominal, title: "Significant?"]
    ])
    |> Kino.VegaLite.new()
  end

  defp interpret_eta_squared(eta) when eta < 0.01, do: "Small effect"
  defp interpret_eta_squared(eta) when eta < 0.06, do: "Medium effect"
  defp interpret_eta_squared(_), do: "Large effect"

  defp interpret_cohens_d(d) when d < 0.2, do: "Negligible effect"
  defp interpret_cohens_d(d) when d < 0.5, do: "Small effect"
  defp interpret_cohens_d(d) when d < 0.8, do: "Medium effect"
  defp interpret_cohens_d(_), do: "Large effect"

  defp format_number(num) when is_float(num) do
    :erlang.float_to_binary(num, decimals: 2)
  end

  defp format_number(num), do: to_string(num)

  defp format_scientific(num) when num < 0.001 do
    :io_lib.format("~.2e", [num]) |> to_string()
  end

  defp format_scientific(num) do
    :erlang.float_to_binary(num, decimals: 5)
  end
end