defmodule Scholar.Metrics.ClusteringTest do
  use Scholar.Case, async: true
  alias Scholar.Metrics.Clustering
  doctest Clustering

  defp x do
    Nx.tensor([
      [0.8157184486105233, -0.01635270455208079, -0.18865907479231095],
      [-0.23329832770253972, -1.0670578022770247, -1.4104723030890731],
      [-1.7011914621798259, -0.5470603319369067, 0.6322584168194634],
      [0.9425204058473541, -1.037556635542662, -0.21666787618406214],
      [-0.3690472295854992, 0.09486339114365348, 0.6961098873520952],
      [-1.6771195334192937, -1.2952988072907938, 0.08604298953131055],
      [-0.15949996798407284, 0.6000955647811537, 0.8663542127627641],
      [-1.7895457106619646, -0.1754378695049237, -0.365205321101012],
      [-1.394493075734099, 0.07372952442240266, 0.050092862319408904],
      [0.7307430281513498, -0.9238089730338562, -1.1574739042650035],
      [0.8915021818008122, 0.3179163863864429, -0.08149978809223538],
      [-3.168762672235742, -0.9546695365525117, -0.9908465332175033],
      [1.4631690777205948, -0.4670129058849814, 0.6461228097855395],
      [0.5622446765584268, -1.3052125865380872, 2.1498777623693526],
      [1.520536826767848, 0.13949842485542288, 0.6555116085852244],
      [0.07897976746203501, -0.6255090568811114, 0.5807450128292253],
      [0.9392107340484556, -0.3814452707719886, 0.30010654905670114],
      [2.2255470282884557, 0.6745085423783157, -1.6413094715372527],
      [-1.1522920682636968, 1.6005950153494766, -0.3385979516219887],
      [0.391502390567214, -0.25562956508471424, 0.12315775252511421]
    ])
  end

  defp y do
    Nx.tensor([2, 2, 2, 2, 0, 2, 1, 1, 2, 2, 2, 2, 2, 1, 1, 0, 1, 2, 1, 2])
  end

  defp y_one do
    Nx.tensor([2, 2, 2, 2, 0, 2, 1, 3, 2, 2, 2, 2, 2, 1, 4, 0, 1, 2, 1, 2])
  end

  test "silhouette_score test" do
    expected = Nx.tensor(-0.17424359382143964)

    assert_all_close(Clustering.silhouette_score(x(), y(), num_clusters: 3), expected)
  end

  test "silhouette_samples test" do
    expected =
      Nx.tensor([
        -0.2700417935848236,
        0.014513458125293255,
        -0.36026713252067566,
        -0.2003783881664276,
        0.4819222092628479,
        -0.22008605301380157,
        -0.526187002658844,
        -0.2967069745063782,
        -0.3657914698123932,
        0.019084271043539047,
        -0.27950319647789,
        -0.03889584541320801,
        -0.3125722408294678,
        -0.3157929480075836,
        -0.23939572274684906,
        0.5096476674079895,
        -0.4214971363544464,
        0.03587301820516586,
        -0.16732777655124664,
        -0.5314690470695496
      ])

    assert_all_close(Clustering.silhouette_samples(x(), y(), num_clusters: 3), expected)
  end

  test "silhouette_score test with one-element clusters" do
    expected = Nx.tensor(-0.298110549250206)

    assert_all_close(Clustering.silhouette_score(x(), y_one(), num_clusters: 5), expected)
  end

  test "silhouette_samples test with one-element clusters" do
    expected =
      Nx.tensor([
        -0.4016384184360504,
        -0.05811696499586105,
        -0.5811521410942078,
        -0.2139229029417038,
        0.45039641857147217,
        -0.5208510756492615,
        -0.5239397287368774,
        0.0,
        -0.7291823625564575,
        0.019084271043539047,
        -0.5117945671081543,
        -0.536014974117279,
        -0.7470640540122986,
        -0.3105389475822449,
        0.0,
        0.4759448170661926,
        -0.6119896173477173,
        -0.2591235637664795,
        -0.37083834409713745,
        -0.5314690470695496
      ])

    assert_all_close(Clustering.silhouette_samples(x(), y_one(), num_clusters: 5), expected)
  end
end
