"""
Copyright 2025 - David Minton.

This file is part of Swiftest.
Swiftest is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License
as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
Swiftest is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along with Swiftest.
If not, see: https://www.gnu.org/licenses.
"""

import os
import tempfile
import unittest
from pathlib import Path

import numpy as np

import swiftest

# Colliion initial conditions taken from Fragmentation_Movie.py in the example directory

collision_type = [
    "disruption_headon",
    "disruption_off_axis",
    "supercatastrophic_headon",
    "supercatastrophic_off_axis",
    "hitandrun_disrupt",
    "hitandrun_pure",
    "merge",
    "merge_spinner",
]

# ----------------------------------------------------------------------------------------------------------------------
# To increase the number of bodies generated in each collision type, decrease the value of the corresponding nfrag_reduction number
# ----------------------------------------------------------------------------------------------------------------------
nfrag_reduction = {
    "disruption_headon": 10.0,
    "disruption_off_axis": 10.0,
    "supercatastrophic_headon": 10.0,
    "supercatastrophic_off_axis": 10.0,
    "hitandrun_disrupt": 10.0,
    "hitandrun_pure": 1.0,
    "merge": 1.0,
    "merge_spinner": 1.0,
}


# These initial conditions were generated by trial and error
names = ["Target", "Projectile"]
pos_vectors = {
    "disruption_headon": [np.array([1.0, -5.0e-05, 0.0]), np.array([1.0, 5.0e-05, 0.0])],
    "disruption_off_axis": [np.array([1.0, -5.0e-05, 0.0]), np.array([1.0, 5.0e-05, 0.0])],
    "supercatastrophic_headon": [np.array([1.0, -5.0e-05, 0.0]), np.array([1.0, 5.0e-05, 0.0])],
    "supercatastrophic_off_axis": [np.array([1.0, -5.0e-05, 0.0]), np.array([1.0, 5.0e-05, 0.0])],
    "hitandrun_disrupt": [np.array([1.0, -4.2e-05, 0.0]), np.array([1.0, 4.2e-05, 0.0])],
    "hitandrun_pure": [np.array([1.0, -4.2e-05, 0.0]), np.array([1.0, 4.2e-05, 0.0])],
    "merge": [np.array([1.0, -5.0e-05, 0.0]), np.array([1.0, 5.0e-05, 0.0])],
    "merge_spinner": [np.array([1.0, -5.0e-05, 0.0]), np.array([1.0, 5.0e-05, 0.0])],
}

vel_vectors = {
    "disruption_headon": [np.array([0.00, 6.280005, 0.0]), np.array([0.00, 3.90, 0.0])],
    "disruption_off_axis": [np.array([0.00, 6.280005, 0.0]), np.array([0.05, 3.90, 0.0])],
    "supercatastrophic_headon": [np.array([0.00, 6.28, 0.0]), np.array([0.00, 4.28, 0.0])],
    "supercatastrophic_off_axis": [np.array([0.00, 6.28, 0.0]), np.array([0.05, 4.28, 0.0])],
    "hitandrun_disrupt": [np.array([0.00, 6.28, 0.0]), np.array([-1.45, -6.28, 0.0])],
    "hitandrun_pure": [np.array([0.00, 6.28, 0.0]), np.array([-1.52, -6.28, 0.0])],
    "merge": [np.array([0.04, 6.28, 0.0]), np.array([0.05, 6.18, 0.0])],
    "merge_spinner": [np.array([0.04, 6.28, 0.0]), np.array([0.05, 6.18, 0.0])],
}

rot_vectors = {
    "disruption_headon": [np.array([0.0, 0.0, 1.0e5]), np.array([0.0, 0.0, -5e5])],
    "disruption_off_axis": [np.array([0.0, 0.0, 2.0e5]), np.array([0.0, 0.0, -1.0e5])],
    "supercatastrophic_headon": [np.array([0.0, 0.0, 1.0e5]), np.array([0.0, 0.0, -5.0e5])],
    "supercatastrophic_off_axis": [np.array([0.0, 0.0, 1.0e5]), np.array([0.0, 0.0, -1.0e4])],
    "hitandrun_disrupt": [np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0e5])],
    "hitandrun_pure": [np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0e5])],
    "merge": [np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, 0.0])],
    "merge_spinner": [np.array([0.0, 0.0, -1.2e6]), np.array([0.0, 0.0, 0.0])],
}

body_Gmass = {
    "disruption_headon": [1e-7, 1e-9],
    "disruption_off_axis": [1e-7, 1e-9],
    "supercatastrophic_headon": [1e-7, 1e-8],
    "supercatastrophic_off_axis": [1e-7, 1e-8],
    "hitandrun_disrupt": [1e-7, 7e-10],
    "hitandrun_pure": [1e-7, 7e-10],
    "merge": [1e-7, 1e-8],
    "merge_spinner": [1e-7, 1e-8],
}

tstop = {
    "disruption_headon": 2.0e-3,
    "disruption_off_axis": 2.0e-3,
    "supercatastrophic_headon": 2.0e-3,
    "supercatastrophic_off_axis": 2.0e-3,
    "hitandrun_disrupt": 2.0e-4,
    "hitandrun_pure": 2.0e-4,
    "merge": 5.0e-3,
    "merge_spinner": 5.0e-3,
}

nfrag_minimum_expected = {
    "disruption_headon": 20,
    "disruption_off_axis": 20,
    "supercatastrophic_headon": 20,
    "supercatastrophic_off_axis": 20,
    "hitandrun_disrupt": 20,
    "hitandrun_pure": 0,
    "merge": 0,
    "merge_spinner": 0,
}

nfrag_maximum_expected = {
    "disruption_headon": 100,
    "disruption_off_axis": 100,
    "supercatastrophic_headon": 100,
    "supercatastrophic_off_axis": 100,
    "hitandrun_disrupt": 100,
    "hitandrun_pure": 0,
    "merge": 0,
    "merge_spinner": 0,
}

expected_regime = {
    "disruption_headon": "Disruption",
    "disruption_off_axis": "Disruption",
    "supercatastrophic_headon": "Supercatastrophic disruption",
    "supercatastrophic_off_axis": "Supercatastrophic disruption",
    "hitandrun_disrupt": "Hit and run",
    "hitandrun_pure": "Hit and run",
    "merge": "Merge",
    "merge_spinner": "Hit and run",
}
expected_outcome = {
    "disruption_headon": "calculation converged",
    "disruption_off_axis": "calculation converged",
    "supercatastrophic_headon": "calculation converged",
    "supercatastrophic_off_axis": "calculation converged",
    "hitandrun_disrupt": "calculation converged",
    "hitandrun_pure": "No new fragments generated",
    "merge": "Merging",
    "merge_spinner": "No new fragments generated",
}

density = 3000 * swiftest.AU2M**3 / swiftest.MSun
GU = swiftest.GMSun * swiftest.YR2S**2 / swiftest.AU2M**3
body_radius = body_Gmass.copy()
for k, v in body_Gmass.items():
    body_radius[k] = [((Gmass / GU) / (4.0 / 3.0 * np.pi * density)) ** (1.0 / 3.0) for Gmass in v]

body_radius["hitandrun_disrupt"] = [7e-6, 3.25e-6]
body_radius["hitandrun_pure"] = [7e-6, 3.25e-6]


class TestSwiftestRestart(unittest.TestCase):
    def setUp(self):
        # Initialize a target and surface for testing
        self.tmpdir = tempfile.TemporaryDirectory()
        self.simdir = self.tmpdir.name

        self.tmpdir2 = tempfile.TemporaryDirectory()
        self.simdir_repeat = self.tmpdir2.name

    def tearDown(self):
        # Clean up temporary directory
        self.tmpdir.cleanup()

    def test_restart_collision_snapshot(self):
        """
        Test that a simulation with collisions restarts.
        """
        sim = swiftest.Simulation(
            simdir=self.simdir,
            integrator="symba",
            dump_cadence=1,
            tstep_out=1 * 365.25,
            tstop=365.25 * 1,
            dt=0.01,
            DU="km",
            TU="d",
            MU2KG=1e15,
        )
        sim.add_solar_system_body(["Mars", "Phobos", "Deimos"])

        # Add Mars c_lm from GMM3 (hardcoded so we don't have to rely on pyshtools being installed)
        c_lm = np.array(
            [
                [
                    [1.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [-8.75021132e-04, 5.90314960e-10, -8.46359039e-05, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [-1.18960349e-05, 3.80429811e-06, -1.59480091e-05, 3.50541857e-05, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [5.12940730e-06, 4.21623326e-06, -9.52539393e-07, 6.45688952e-06, 3.09718541e-07, 0.00000000e00, 0.00000000e00],
                    [
                        -1.72668240e-06,
                        4.83798784e-07,
                        -4.29814806e-06,
                        3.31265594e-06,
                        -4.64037477e-06,
                        -4.44982150e-06,
                        0.00000000e00,
                    ],
                    [
                        1.34648814e-06,
                        1.80236722e-06,
                        8.61611560e-07,
                        9.55649930e-07,
                        1.00849633e-06,
                        1.65767871e-06,
                        2.76208184e-06,
                    ],
                ],
                [
                    [0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [0.00000000e00, -4.94336174e-11, 4.89346259e-05, 0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [0.00000000e00, 2.51773217e-05, 8.36151763e-06, 2.55745706e-05, 0.00000000e00, 0.00000000e00, 0.00000000e00],
                    [
                        0.00000000e00,
                        3.76338480e-06,
                        -8.98081277e-06,
                        -1.93196208e-07,
                        -1.28730468e-05,
                        0.00000000e00,
                        0.00000000e00,
                    ],
                    [
                        0.00000000e00,
                        2.12322092e-06,
                        -1.16588966e-06,
                        2.71791890e-07,
                        -3.38212383e-06,
                        3.77983077e-06,
                        0.00000000e00,
                    ],
                    [
                        0.00000000e00,
                        -1.51845907e-06,
                        1.46912740e-06,
                        3.33031890e-07,
                        2.63880041e-06,
                        1.62290102e-06,
                        8.21838201e-07,
                    ],
                ],
            ]
        )

        sim.modify_body(name="Mars", c_lm=c_lm)

        # hardcoded impactors to prevent use of my git repo

        names = np.array(["100", "101", "102", "103", "104", "105", "106", "107", "108", "109", "598"])
        radius = np.array(
            [
                0.35705834,
                0.07421946,
                0.14916194,
                0.13296731,
                0.12915766,
                0.41544253,
                0.46550525,
                0.17445848,
                0.41888955,
                0.45045171,
                0.49215104,
            ]
        )
        mass = 10 * np.array(
            [
                2.86020662e-04,
                2.56881727e-06,
                2.08522999e-05,
                1.47711587e-05,
                1.35375579e-05,
                4.50518599e-04,
                6.33802015e-04,
                3.33623100e-05,
                4.61826081e-04,
                5.74281193e-04,
                7.48988414e-04,
            ]
        )
        rh = np.array(
            [
                [-22124.94631599, -7600.61943767, -1953.85576394],
                [-22099.04261316, -7594.5053911, -1955.35554083],
                [-22122.04312455, -7615.95767634, -1955.17148458],
                [-22083.30057286, -7614.06013637, -1953.85760781],
                [-22107.9145528, -7592.30630884, -1950.69068123],
                [-22124.83908881, -7600.49807585, -1961.07379115],
                [-22091.70625466, -7601.39044401, -1943.88804141],
                [-22091.96202388, -7598.98000554, -1952.81119451],
                [-22098.41518385, -7590.14851984, -1954.86866938],
                [-22117.17018901, -7620.37817046, -1965.44576821],
                [-22104.19743262, -7621.76082186, -1942.37784931],
            ]
        )
        vh = 0.503 * np.array(
            [
                [36796.54983816, -110084.06529807, 2987.13829165],
                [38083.48622951, -110102.85015476, 2921.36450056],
                [36869.30640377, -110484.65220524, 2935.57333344],
                [38278.23679524, -110732.77923, 2990.69525294],
                [37624.67760602, -109538.53678746, 3189.04717925],
                [36792.03452032, -109813.24269459, 2613.69041106],
                [38540.20902575, -110535.5996834, 3651.11554834],
                [38768.73186052, -109936.55975385, 3071.4989985],
                [38141.5822543, -109706.74112105, 2931.3008305],
                [36896.44746876, -111001.2491618, 2413.7897524],
                [37671.8566089, -110761.29152044, 3217.24051514],
            ]
        )

        sim.add_body(name=names, radius=radius, mass=mass, rh=rh, vh=vh)

        try:
            sim.run()
        except Exception as e:
            self.fail(f"Failed initial run with Exception: {e}")

        # restart run

        sim_restart = swiftest.Simulation(simdir=self.simdir, read_data=True, param_file="param.restart.in", tstop=365.25 * 10)

        try:
            sim_restart.run()
        except Exception as e:
            self.fail(f"Failed restart with Exception: {e}")

        self.assertTrue((Path(self.simdir) / "param.00000000000000365250.in").is_file())
        self.assertEqual(sim_restart.data.time.size, 11)
        return

    def test_restart_accurate(self):
        """
        Test that a restarted run gives the same outputs as a full run.

        """
        # Define some common arguments
        run_args = {"tstop": 10, "dt": 0.01, "tstep_out": 1, "dump_cadence": 1, "compute_conservation_values": True}

        integrators = ["whm", "helio", "rmvs", "symba"]
        bodies = ["Sun", "Mercury", "Venus", "Earth", "Mars"]
        restart_iout = 500
        param_restart = f"param.{restart_iout:020d}.in"
        for i in integrators:
            sim = swiftest.Simulation(simdir=self.simdir, integrator=i, **run_args, clean=True)
            sim.add_solar_system_body(bodies)
            try:
                sim.run()
            except Exception as e:
                self.fail(f"Failed initial run with Exception: {e}")

            # repeat run with exact same parameters
            sim_repeat = swiftest.Simulation(simdir=self.simdir_repeat, integrator=i, **run_args, clean=True)
            sim_repeat.add_solar_system_body(bodies)

            try:
                sim_repeat.run()
            except Exception as e:
                self.fail(f"Failed repeat run with Exception: {e}")

            # Check that the output files are the same
            for var in sim.data.data_vars:
                if sim.data[var].dtype.type != np.str_ and np.isnan(sim.data[var].values).any():
                    idx = np.where(~np.isnan(sim.data[var].values))
                    self.assertTrue(
                        (sim.data[var].values[idx] == sim_repeat.data[var].values[idx]).all(),
                        f"{i}: {var} values are not equal\n\n{sim.data[var].values[idx]}\n\nREPEAT\n{sim_repeat.data[var].values[idx]}",
                    )
                else:
                    self.assertTrue(
                        (sim.data[var].values == sim_repeat.data[var].values).all(),
                        f"{var} values are not equal\n\n{sim.data[var].values}\n\nREPEAT\n{sim_repeat.data[var].values}",
                    )

            # restarted run (from halfway mark in this case)
            sim_restart = swiftest.Simulation(
                simdir=self.simdir, integrator=i, read_data=True, param_file=param_restart, compute_conservation_values=True
            )

            try:
                sim_restart.run()
            except Exception as e:
                self.fail(f"Failed restart run with Exception: {e}")

            # check that the output files are the same
            for var in sim.data.data_vars:
                if sim.data[var].dtype.type != np.str_ and np.isnan(sim.data[var].values).any():
                    idx = np.where(~np.isnan(sim.data[var].values))
                    self.assertTrue(
                        (sim.data[var].values[idx] == sim_restart.data[var].values[idx]).all(),
                        f"{i}: {var} values are not equal\n\n{sim.data[var].values[idx]}\n\nRESTART\n{sim_restart.data[var].values[idx]}",
                    )
                else:
                    self.assertTrue(
                        (sim.data[var].values == sim_restart.data[var].values).all(),
                        f"{i}: {var} values are not equal\n\n{sim.data[var].values}\n\nRESTART\n{sim_restart.data[var].values}",
                    )
        return

    def test_restart_accurate_collision(self):
        """
        Test that a restarted run with collisions gives the same outputs as a full run.

        Collision outcomes taken from test_fraggle.py

        """
        seed = [974899978, -2041855733, 14415898, 615945619, 1818808148, 462758063, 762516551, 1276432020]

        for style in collision_type:
            sim = swiftest.Simulation(simdir=self.simdir, rotation=True, compute_conservation_values=True, seed=seed, clean=True)
            sim.add_solar_system_body("Sun")
            sim.add_body(
                name=names,
                Gmass=body_Gmass[style],
                radius=body_radius[style],
                rh=pos_vectors[style],
                vh=vel_vectors[style],
                rot=rot_vectors[style],
            )

            # Set fragmentation parameters
            minimum_fragment_gmass = 0.01 * body_Gmass[style][1]
            gmtiny = 0.50 * body_Gmass[style][1]
            sim.set_parameter(
                encounter_save="both",
                gmtiny=gmtiny,
                minimum_fragment_gmass=minimum_fragment_gmass,
                nfrag_reduction=nfrag_reduction[style],
            )
            try:
                sim.run(dt=tstop[style] / 4, tstop=tstop[style], istep_out=1, dump_cadence=2)
            except Exception as e:
                self.fail(f"Failed initial run with Exception: {e}")

            # repeat run with exact same parameters
            sim_repeat = swiftest.Simulation(
                simdir=self.simdir_repeat, rotation=True, compute_conservation_values=True, seed=seed, clean=True
            )
            sim_repeat.add_solar_system_body("Sun")
            sim_repeat.add_body(
                name=names,
                Gmass=body_Gmass[style],
                radius=body_radius[style],
                rh=pos_vectors[style],
                vh=vel_vectors[style],
                rot=rot_vectors[style],
            )

            # Set fragmentation parameters
            sim_repeat.set_parameter(
                encounter_save="both",
                gmtiny=gmtiny,
                minimum_fragment_gmass=minimum_fragment_gmass,
                nfrag_reduction=nfrag_reduction[style],
            )
            try:
                sim_repeat.run(dt=tstop[style] / 4, tstop=tstop[style], istep_out=1, dump_cadence=2)
            except Exception as e:
                self.fail(f"Failed repeat run with Exception: {e}")

                # check that the output files are the same

            for var in sim.data.data_vars:
                if sim.data[var].dtype.type != np.str_ and np.isnan(sim.data[var].values).any():
                    idx = np.where(~np.isnan(sim.data[var].values))
                    self.assertTrue(
                        (sim.data[var].values[idx] == sim_repeat.data[var].values[idx]).all(),
                        f"{var} values are not equal\n\n{sim.data[var].values[idx]}\n\nREPEAT\n{sim_repeat.data[var].values[idx]}",
                    )
                else:
                    self.assertTrue(
                        (sim.data[var].values == sim_repeat.data[var].values).all(),
                        f"{var} values are not equal\n\n{sim.data[var].values}\n\nREPEAT\n{sim_repeat.data[var].values}",
                    )

            # restarted run (from the halfway mark in this case)

            restart_iout = 2
            param_restart = f"param.{restart_iout:020d}.in"

            sim.set_parameter(param_file=param_restart)
            try:
                sim.run()
            except Exception as e:
                self.fail(f"Failed restart run with Exception: {e}")

                # check that the output files are the same

            for var in sim.data.data_vars:
                if sim.data[var].dtype.type != np.str_ and np.isnan(sim.data[var].values).any():
                    idx = np.where(~np.isnan(sim.data[var].values))
                    self.assertTrue(
                        (sim.data[var].values[idx] == sim_repeat.data[var].values[idx]).all(),
                        f"{var} values are not equal\n\n{sim.data[var].values[idx]}\n\nRESTART\n{sim_repeat.data[var].values[idx]}",
                    )
                else:
                    self.assertTrue(
                        (sim.data[var].values == sim_repeat.data[var].values).all(),
                        f"{var} values are not equal\n\n{sim.data[var].values}\n\nRESTART\n{sim_repeat.data[var].values}",
                    )

        return


if __name__ == "__main__":
    os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"
    unittest.main()
