import os
import re
import subprocess
from pathlib import Path
from typing import List

import pytest


class Env:
    def __init__(self, ref_filename: str, result_filename: str, input_filename: str, options: str):
        self.test_path = Path(__file__).resolve().parent
        self.ref_filepath = Path.joinpath(self.test_path, "references", ref_filename)
        self.result_filepath = Path.joinpath(self.test_path, "results", result_filename)
        self.input_filepath = Path.joinpath(self.test_path, "data", input_filename)
        self.command: str = f"sum_dirac_dfcoef -i {self.input_filepath} -o {self.result_filepath} {options}"


def get_output_list(filepath: Path) -> List[List[str]]:
    return [re.split(" +", line.rstrip("\n")) for line in list(filter(lambda val: val != "", open(filepath).read().splitlines()))]


@pytest.mark.parametrize(
    "ref_filename, result_filename, input_filename, options",
    # fmt: off
    [
        ("ref.Ar.compress.out"                      , "result.Ar.compress.out"                      , "Ar_Ar.out"         , "-d 15 -c"),
        ("ref.Ar.no_sort.compress.out"              , "result.Ar.no_sort.compress.out"              , "Ar_Ar.out"         , "-d 15 --no-sort -c"),
        ("ref.N2.compress.out"                      , "result.N2.compress.out"                      , "N2_N2.out"         , "-d 15 -c"),
        ("ref.N2.compress.positronic.out"           , "result.N2.compress.positronic.out"           , "N2_N2.out"         , "-d 15 -pc"),
        ("ref.N2.compress.all.out"                  , "result.N2.compress.all.out"                  , "N2_N2.out"         , "-d 15 -ac"),
        ("ref.N2.no_sort.compress.out"              , "result.N2.no_sort.compress.out"              , "N2_N2.out"         , "-d 15 --no-sort -c"),
        ("ref.N2.no_sort.compress.positronic.out"   , "result.N2.no_sort.compress.positronic.out"   , "N2_N2.out"         , "-d 15 --no-sort -pc"),
        ("ref.N2.no_sort.compress.all.out"          , "result.N2.no_sort.compress.all.out"          , "N2_N2.out"         , "-d 15 --no-sort -ac"),
        ("ref.uo2.compress.out"                     , "result.uo2.compress.out"                     , "x2c_uo2_238.out"   , "-d 15 -c"),
        ("ref.uo2.no_sort.compress.out"             , "result.uo2.no_sort.compress.out"             , "x2c_uo2_238.out"   , "-d 15 --no-sort -c"),
        ("ref.ucl4.compress.out"                    , "result.ucl4.compress.out"                    , "x2c_ucl4.out"      , "-d 15 -c"),
        ("ref.ucl4.no_sort.compress.out"            , "result.ucl4.no_sort.compress.out"            , "x2c_ucl4.out"      , "-d 15 --no-sort -c"),
        ("ref.Cm3+_phen.compress.out"               , "result.Cm3+_phen.compress.out"               , "x2c_Cm3+_phen.out" , "-d 15 -c"),
    ]
    # fmt: on
)
def test_sum_dirac_dfcoeff_compress(ref_filename: str, result_filename: str, input_filename: str, options: str):
    env = Env(ref_filename, result_filename, input_filename, options)
    os.chdir(env.test_path)
    print(f"{env.test_path} test start...\ncommand: {env.command}")
    subprocess.run(env.command.split(), encoding="utf-8", check=True)

    ref_list: List[List[str]] = get_output_list(env.ref_filepath)
    result_list: List[List[str]] = get_output_list(env.result_filepath)
    # File should have the same number of lines
    assert len(ref_list) == len(result_list), f"Number of lines in {ref_filename}(={len(ref_list)}) and {result_filename}(={len(result_list)}) are different."
    threshold: float = 1e-10
    for line_idx, (ref, out) in enumerate(zip(ref_list, result_list)):
        # 1st and 2nd lines have header information about eigenvalues
        # (e.g.) electron_num 18 E1g 1..33 E1u 1..33
        #        E1g closed 6 open 0 virtual 60 E1u closed 12 open 0 virtual 54
        if line_idx < 2:
            assert ref == out
            continue
        # ref[0]: irrep, ref[1]: energy order index in the irrep, ref[2]: energy, ref[3:]: Symmetry value and coefficient
        # (e.g.) E1u 19 -8.8824415703374 B3uUpx 49.999172476298732 B2uUpy 49.999172476298732
        assert ref[0] == out[0], f"irrep in line {line_idx} of {ref_filename} and {result_filename} are different."
        assert ref[1] == out[1], f"Energy order index in line {line_idx} of {ref_filename} and {result_filename} are different."
        assert float(ref[2]) == pytest.approx(float(out[2]), abs=threshold), f"Energy in line {line_idx} of {ref_filename} and {result_filename} are different."
        for idx, (ref_val, out_val) in enumerate(zip(ref[3:], out[3:])):
            if idx % 2 == 0:
                assert ref_val == out_val, f"Symmetry value in line {line_idx} of {ref_filename} and {result_filename} are different."
            else:
                assert float(ref_val) == pytest.approx(
                    float(out_val), abs=threshold
                ), f"Contribution of the AO in the MO in line {line_idx} of {ref_filename} and {result_filename} are different."


@pytest.mark.parametrize(
    "ref_filename, result_filename, input_filename, options",
    # fmt: off
    [
        ("ref.Ar.out"                       , "result.Ar.out"                       , "Ar_Ar.out"                       , "-d 15"),
        ("ref.Ar.no_sort.out"               , "result.Ar.no_sort.out"               , "Ar_Ar.out"                       , "-d 15 --no-sort"),
        ("ref.N2.out"                       , "result.N2.out"                       , "N2_N2.out"                       , "-d 15"),
        ("ref.N2.positronic.out"            , "result.N2.positronic.out"            , "N2_N2.out"                       , "-d 15 -p"),
        ("ref.N2.all.out"                   , "result.N2.all.out"                   , "N2_N2.out"                       , "-d 15 -a"),
        ("ref.N2.no_sort.out"               , "result.N2.no_sort.out"               , "N2_N2.out"                       , "-d 15 --no-sort"),
        ("ref.N2.no_sort.positronic.out"    , "result.N2.no_sort.positronic.out"    , "N2_N2.out"                       , "-d 15 --no-sort -p"),
        ("ref.N2.no_sort.all.out"           , "result.N2.no_sort.all.out"           , "N2_N2.out"                       , "-d 15 --no-sort -a"),
        ("ref.uo2.special.out"              , "result.uo2.special.out"              , "special_exit_condition_UO2.out"  , "-d 15"),
        ("ref.uo2.out"                      , "result.uo2.out"                      , "x2c_uo2_238.out"                 , "-d 15"),
        ("ref.uo2.no_sort.out"              , "result.uo2.no_sort.out"              , "x2c_uo2_238.out"                 , "-d 15 --no-sort"),
        ("ref.ucl4.out"                     , "result.ucl4.out"                     , "x2c_ucl4.out"                    , "-d 15"),
        ("ref.ucl4.no_sort.out"             , "result.ucl4.no_sort.out"             , "x2c_ucl4.out"                    , "-d 15 --no-sort"),
        ("ref.Cm3+_phen.out"                , "result.Cm3+_phen.out"                , "x2c_Cm3+_phen.out"               , "-d 15"),
    ]
    # fmt: on
)
def test_sum_dirac_dfcoeff(ref_filename: str, result_filename: str, input_filename: str, options: str):
    env = Env(ref_filename, result_filename, input_filename, options)
    os.chdir(env.test_path)
    print(f"{env.test_path} test start...\ncommand: {env.command}")
    subprocess.run(env.command.split(), encoding="utf-8", check=True)

    ref_list: List[List[str]] = get_output_list(env.ref_filepath)
    result_list: List[List[str]] = get_output_list(env.result_filepath)
    # File should have the same number of lines
    assert len(ref_list) == len(result_list), f"Number of lines in {ref_filename}(={len(ref_list)}) and {result_filename}(={len(result_list)}) are different."
    threshold: float = 1e-10
    for line_idx, (ref, out) in enumerate(zip(ref_list, result_list)):
        # 1st and 2nd lines have header information about eigenvalues
        # (e.g.) electron_num 18 E1g 1..33 E1u 1..33
        #        E1g closed 6 open 0 virtual 60 E1u closed 12 open 0 virtual 54
        if line_idx < 2:
            assert ref == out
            continue
        if len(ref) < 2 or len(out) < 2:
            continue
        if "%" in ref[-1]:
            ref_value = float(ref[-2])
            out_value = float(out[-2])
            ref_list_str = " ".join(ref[:-2])
            out_list_str = " ".join(out[:-2])
        else:
            ref_value = float(ref[-1])
            out_value = float(out[-1])
            ref_list_str = " ".join(ref[:-1])
            out_list_str = " ".join(out[:-1])
        assert ref_list_str == out_list_str, f"line {line_idx}: {ref_list_str} != {out_list_str}\nref: {ref_list[line_idx]}\nout:{result_list[line_idx]}"
        assert ref_value == pytest.approx(out_value, abs=threshold), f"line {line_idx}: {ref_value} != {out_value}\nref: {ref_list[line_idx]}\nout:{result_list[line_idx]}"


def test_version_option():
    command = "sum_dirac_dfcoef -v"
    p = subprocess.run(command.split(), encoding="utf-8", check=True, stdout=subprocess.PIPE)
    # p.stdout = x.y.z\n
    out_version_str = p.stdout.rstrip("\n")
    about_file = Path(__file__).resolve().parent.parent / "src/sum_dirac_dfcoef/__about__.py"
    # readline() => __version__ = "x.y.z"
    # split()[-1] => "x.y.z"
    # replace() => x.y.z
    ref_version_str = open(about_file).readline().split()[-1].replace('"', "")
    assert out_version_str == ref_version_str


def test_help_option():
    command = "sum_dirac_dfcoef -h"
    subprocess.run(command.split(), check=True)
