"""Tests of the templates module."""

from itertools import chain
from pathlib import Path

import pytest

from py21cmfast import _templates as tmpl
from py21cmfast.wrapper.inputs import CosmoTables, InputParameters, InputStruct

_TEMPLATES = tmpl.list_templates()
_ALL_ALIASES = list(chain.from_iterable(t["aliases"] for t in _TEMPLATES))


class TestListTemplates:
    """Tests of the list_templates function."""

    def setup_class(self):
        """Set up the class."""
        self.templates = _TEMPLATES

    def test_output_form(self):
        """Test that list_templates() returns the correct format of list."""
        assert isinstance(self.templates, list)
        assert all(isinstance(x, dict) for x in self.templates)

    def test_existence_of_files(self):
        """Ensure that all template files exist."""
        assert all((tmpl.TEMPLATE_PATH / t["file"]).exists() for t in self.templates)

    def test_uniqueness_of_names(self):
        """Test that all names in the templates are unique."""
        allnames = [t["name"] for t in self.templates]
        assert len(set(allnames)) == len(allnames)

    def test_uniqueness_of_aliases(self):
        """Test that all aliases in the templates are unique."""
        assert len(set(_ALL_ALIASES)) == len(_ALL_ALIASES)


class TestLoadTemplateFile:
    """Tests of the load_template_file function."""

    @pytest.mark.parametrize("template", _ALL_ALIASES)
    def test_load_str(self, template: str):
        """Test that all templates only have correct InputStruct keys in them."""
        out = tmpl.load_template_file(template)
        possible = set(InputStruct._subclasses.keys())
        assert all(key in possible for key in out)

    def test_load_path(self, tmp_path: Path):
        """Test that loading from an explicit TOML path works."""
        pth = tmp_path / "tmp.toml"
        with pth.open("w") as fl:
            fl.write(
                """
                [SimulationOptions]
                BOX_LEN = 100
                """
            )

        out = tmpl.load_template_file(pth)
        assert out["SimulationOptions"]["BOX_LEN"] == 100

    def test_non_existent(self):
        """Test that loading a non-existent template errors."""
        with pytest.raises(ValueError, match="Template VIRIDISSUCKS not found on-disk"):
            tmpl.load_template_file("VIRIDISSUCKS")


class TestCreateParamsFromTemplate:
    """Tests of the create_params_from_template function."""

    def test_string_template_equiv_to_list_of_one_string(self):
        """Test that loading a string is the same as loading a len-1 list."""
        inp1 = tmpl.create_params_from_template("simple")
        inp2 = tmpl.create_params_from_template(["simple"])
        assert inp1 == inp2

    @pytest.mark.parametrize("template", _ALL_ALIASES)
    def test_each_template_works(self, template: str):
        """Test that each template creates a dict of input structs."""
        dct = tmpl.create_params_from_template(template)
        cosmo_tables = dct.pop("cosmo_tables")
        assert all(isinstance(v, InputStruct) for v in dct.values())
        assert isinstance(cosmo_tables, CosmoTables)

    def test_adding_kwargs(self):
        """Test that also specifying loose params does set those params."""
        dct = tmpl.create_params_from_template("simple", HII_DIM=1000)
        assert dct["simulation_options"].HII_DIM == 1000


class TestWriteTemplate:
    """Tests of the write_template function."""

    def test_defaults_minimal(self, tmp_path):
        """Test that writing a minimal-mode default file has nothing in it."""
        inputs = InputParameters.from_template("default", random_seed=1)
        pth = tmp_path / "tmp.toml"
        tmpl.write_template(inputs, pth, mode="minimal")

        with pth.open("r") as fl:
            lines = [line for line in fl.readlines() if line]

        assert len(lines) == 2
        assert "This file was generated by py21cmfast" in lines[0]
        assert "Created on" in lines[1]

    @pytest.mark.parametrize("template", _ALL_ALIASES)
    @pytest.mark.parametrize("mode", ["full", "minimal"])
    def test_roundtrip(self, template, tmp_path: Path, mode: str):
        """Test that writing then reading a template gives the same answer."""
        # Some templates require running CLASS. We set K_MAX_FOR_CLASS to be small so the test won't take too long
        inputs = InputParameters.from_template(
            template, random_seed=1, K_MAX_FOR_CLASS=1.0
        )
        pth = tmp_path / "tmp.toml"
        tmpl.write_template(inputs, pth, mode=mode)

        new = InputParameters.from_template(pth, random_seed=1, K_MAX_FOR_CLASS=1.0)
        assert new == inputs
