defmodule Scholar.Manifold.TSNETest do
  use Scholar.Case, async: true
  alias Scholar.Manifold.TSNE
  doctest TSNE

  defp key do
    Nx.Random.key(42)
  end

  defp x do
    Nx.tensor([
      [
        -8.73180761,
        -17.31599354,
        -10.76061306,
        44.89164949,
        3.12358324,
        -12.12349236
      ],
      [
        7.64626495,
        18.55374334,
        3.87725984,
        -6.99299854,
        -9.36734843,
        4.55923917
      ],
      [
        20.09781323,
        8.68550196,
        8.43256455,
        -19.15829538,
        31.40440792,
        6.06427556
      ],
      [
        10.63844377,
        12.75106034,
        1.43716746,
        22.80253916,
        -30.5268915,
        9.39948666
      ],
      [
        4.86846911,
        -28.12055241,
        -6.08674611,
        17.03477114,
        6.65227597,
        -32.9099819
      ],
      [
        -6.85085101,
        -16.87612764,
        -11.54632153,
        -51.06982144,
        -6.74068325,
        8.01699475
      ],
      [
        5.54656739,
        38.3957026,
        24.36609113,
        6.97242576,
        30.16740117,
        -28.33044809
      ],
      [
        -9.46902188,
        -28.02162129,
        37.34167044,
        -5.5330473,
        2.98759518,
        -19.68337166
      ],
      [
        8.7946142,
        0.15558464,
        0.47478064,
        26.20112788,
        47.17380434,
        19.2504862
      ],
      [
        -9.549788,
        7.08735882,
        30.68402683,
        31.23933105,
        4.33932039,
        11.94982724
      ]
    ])
  end

  defp x_wide do
    Nx.tensor([[1, 2, 3, 1], [2, 56, 2, 4]])
  end

  test "all default params" do
    embedding = TSNE.fit(x(), key: key())

    expected =
      Nx.tensor([
        [-9.765890121459961, 0.9379611015319824],
        [-11.129755973815918, 0.9517970085144043],
        [26.718856811523438, -19.29998779296875],
        [-24.4569091796875, 4.920958518981934],
        [-1.2708511352539062, -5.8597612380981445],
        [125.60948181152344, -28.224010467529297],
        [-34.028900146484375, 9.524750709533691],
        [-39.59464645385742, 18.714412689208984],
        [-8.00776195526123, 14.913147926330566],
        [-24.073219299316406, 3.4196181297302246]
      ])

    assert_all_close(embedding, expected)
  end

  test "test for wide matrix" do
    embedding = TSNE.fit(x_wide(), key: key())

    expected =
      Nx.tensor([
        [3679.796875, 0.0],
        [-3679.796875, 0.0]
      ])

    assert_all_close(embedding, expected)
  end

  test "non-default num_components" do
    embedding = TSNE.fit(x(), num_components: 3, key: key())

    expected =
      Nx.tensor([
        [3.5391688346862793, 0.6509681344032288, -2.010162830352783],
        [-1.2239890098571777, -1.5295288562774658, 10.722641944885254],
        [-27.263996124267578, -14.321371078491211, 0.9299386739730835],
        [2.164340019226074, -3.805640935897827, -0.5478832721710205],
        [-1.0675201416015625, -3.0648369789123535, 4.745659351348877],
        [-4.2959489822387695, -1.9941456317901611, -1.3964862823486328],
        [1.1918673515319824, 4.0927414894104, -3.782269239425659],
        [24.376800537109375, 8.162726402282715, 1.711245059967041],
        [0.6383852958679199, 1.2189958095550537, -3.991114377975464],
        [1.93916654586792, 10.589120864868164, -6.380754470825195]
      ])

    assert_all_close(embedding, expected)
  end

  test "non-default perplexity" do
    embedding = TSNE.fit(x(), perplexity: 5, key: key())

    expected =
      Nx.tensor([
        [-92.12468719482422, 14.195682525634766],
        [126.61585235595703, -35.86635971069336],
        [-73.21896362304688, -5.0058393478393555],
        [66.74797058105469, -13.232538223266602],
        [-36.686031341552734, 1.8704216480255127],
        [101.69606018066406, 44.581092834472656],
        [3.147512435913086, -51.27547073364258],
        [-65.03691101074219, 3.034095287322998],
        [89.28443145751953, -9.13759708404541],
        [-120.42665100097656, 50.83637619018555]
      ])

    assert_all_close(embedding, expected)
  end

  test "non-default learning_rate" do
    embedding = TSNE.fit(x(), learning_rate: 100, key: key())

    expected =
      Nx.tensor([
        [-1.3461402654647827, 7.358520984649658],
        [2.0062308311462402, -2.3735902309417725],
        [-12.489317893981934, 15.401891708374023],
        [4.150882720947266, -1.4218299388885498],
        [7.043045997619629, -2.9974515438079834],
        [-7.473917484283447, 4.841724395751953],
        [-8.264013290405273, 5.41910982131958],
        [0.048639118671417236, -1.313403606414795],
        [2.767026662826538, -1.5423762798309326],
        [13.557697296142578, -23.37268829345703]
      ])

    assert_all_close(embedding, expected)
  end

  test "non-default max_iter" do
    embedding = TSNE.fit(x(), num_iters: 333, key: key())

    expected =
      Nx.tensor([
        [-30.11622428894043, 3.6230711936950684],
        [-0.13294053077697754, -3.0624148845672607],
        [64.19924926757812, 2.6656923294067383],
        [-7.7999677658081055, -0.43307507038116455],
        [-37.81877899169922, 9.290597915649414],
        [120.86727142333984, 24.378297805786133],
        [-119.47084045410156, -27.3742618560791],
        [12.207932472229004, -2.9396910667419434],
        [-1.9976780414581299, -5.294439792633057],
        [0.06211590766906738, -0.8550639152526855]
      ])

    assert_all_close(embedding, expected)
  end

  test "non-default init" do
    embedding = TSNE.fit(x(), init: :random, key: key())

    expected =
      Nx.tensor([
        [-156.66343688964844, -34.973445892333984],
        [78.5311050415039, -101.2839126586914],
        [75.17900848388672, 34.4267692565918],
        [-43.72159194946289, -102.33932495117188],
        [153.0662078857422, 101.59364318847656],
        [84.43276977539062, 0.04915308952331543],
        [-44.736968994140625, -35.98069381713867],
        [-79.83869171142578, 370.6204528808594],
        [-32.189876556396484, -120.83247375488281],
        [-34.056640625, -111.28446960449219]
      ])

    assert_all_close(embedding, expected)
  end

  test "non-default metric" do
    embedding = TSNE.fit(x(), key: key(), metric: :chebyshev)

    expected =
      Nx.tensor([
        [9.365217208862305, 26.647315979003906],
        [4.629116058349609, -2.9455509185791016],
        [49.64634323120117, -11.4344482421875],
        [8.727602005004883, 25.68651008605957],
        [-33.325767517089844, -12.945419311523438],
        [2.2205429077148438, 34.78094482421875],
        [42.87562561035156, 35.790077209472656],
        [-35.2755241394043, -112.85540008544922],
        [-31.687278747558594, -16.406055450439453],
        [-17.169509887695312, 33.69422149658203]
      ])

    assert_all_close(embedding, expected)
  end

  test "non-default exaggeration" do
    embedding = TSNE.fit(x(), key: key(), exaggeration: 4.5)

    expected =
      Nx.tensor([
        [0.24389100074768066, 85.93327331542969],
        [-10.559154510498047, -82.89259338378906],
        [2.398576021194458, -8.886427879333496],
        [-6.7537055015563965, -13.02280330657959],
        [-9.773503303527832, -2.8710267543792725],
        [76.82627868652344, -76.25054168701172],
        [6.6305975914001465, 65.817138671875],
        [4.859842300415039, -8.069650650024414],
        [-61.29663848876953, 86.92620086669922],
        [-2.575655698776245, -46.68356704711914]
      ])

    assert_all_close(embedding, expected)
  end
end
