# Copyright 2023-2026 The PULGON Project Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.


import itertools
import logging
from typing import List, Tuple

import numpy as np
import sympy
from sympy import symbols


def sym_inverse_eye(n: int) -> sympy.Matrix:
    A = sympy.zeros(n)
    for ii in range(n):
        A[ii, n - ii - 1] = 1
    return A


def line_group_sympy_withparities(
    DictParams: dict, symprec: float = 1e-6
) -> Tuple[list, list, list]:
    family = DictParams["family"]

    if family == 2:
        qpoint = DictParams["qpoints"]
        nrot = DictParams["nrot"]
        a = DictParams["a"]
        order = DictParams["order"]
        k1, m1, n, piH = symbols("k1 m1 n piH")
        func0 = sympy.Matrix(
            [
                1,
                sympy.exp(1j * k1 * a),
                piH * sympy.exp(1j * m1 * sympy.pi / n),
            ]
        )
        func1 = [
            sympy.Matrix([[1, 0], [0, 1]]),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * k1 * a),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(-1j * k1 * a),
                    ],
                ]
            ),
            sympy.Matrix(
                [
                    [
                        0,
                        sympy.exp(1j * m1 * 2 * sympy.pi / n),
                    ],
                    [
                        1,
                        0,
                    ],
                ]
            ),
        ]
        func = [func0, func1]
        qps_value = [qpoint]
        m1_value = list(range(-int(nrot / 2) + 1, int(nrot / 2) + 1))

        def value_fc(
            fc: list,
            tmp_k1: float,
            tmp_m1: float,
            tmp_piH: int,
            nrot: int,
            order: List[List[int]],
        ) -> list:
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0
                tmp1 = tmp0.evalf(
                    subs={k1: tmp_k1, m1: tmp_m1, n: nrot, piH: tmp_piH}
                )
                res.append(tmp1)
            return res

        paras_km = list(itertools.product(*[qps_value, m1_value]))
        paras_symbol = [k1, m1, piH]
        characters, paras_values = [], []
        for ii, paras_value in enumerate(paras_km):
            tmp_k1, tmp_m1 = paras_value
            if np.isclose(tmp_k1, 0, atol=symprec) or np.isclose(
                tmp_k1, np.pi / a, atol=symprec
            ):
                idx_fc = 0
                fc = func[idx_fc]
                for tmp_piH in [-1, 1]:
                    res = value_fc(fc, tmp_k1, tmp_m1, tmp_piH, nrot, order)
                    characters.append(np.array(res).astype(np.complex128))
                    paras_values.append([tmp_k1, tmp_m1, tmp_piH])
            else:
                idx_fc = 1
                fc = func[idx_fc]
                tmp_piH = 0
                res = value_fc(fc, tmp_k1, tmp_m1, tmp_piH, nrot, order)
                characters.append(np.array(res).astype(np.complex128))
                paras_values.append([tmp_k1, tmp_m1, tmp_piH])

    elif family == 4:
        qpoint = DictParams["qpoints"]
        nrot = DictParams["nrot"]
        a = DictParams["a"]
        order = DictParams["order"]
        k1, m1, n, piH = symbols("k1 m1 n piH")
        func0 = sympy.Matrix(
            [
                1,
                sympy.exp(1j * m1 * sympy.pi / n),
                sympy.exp(1j * m1 * 2 * sympy.pi / n),
                piH,
            ]
        )
        func1 = [
            sympy.Matrix([[1, 0], [0, 1]]),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * (m1 * sympy.pi / n + k1 * a / 2)),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(1j * (m1 * sympy.pi / n - k1 * a / 2)),
                    ],
                ]
            ),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * m1 * 2 * sympy.pi / n),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(1j * m1 * 2 * sympy.pi / n),
                    ],
                ]
            ),
            sympy.Matrix([[0, 1], [1, 0]]),
        ]
        func = [func0, func1]
        qps_value = [qpoint]
        m1_value = list(range(-int(nrot) + 1, int(nrot) + 1))

        def value_fc(
            fc: list,
            tmp_k1: float,
            tmp_m1: float,
            tmp_piH: int,
            nrot: int,
            order: List[List[int]],
        ) -> list:
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0
                tmp1 = tmp0.evalf(
                    subs={k1: tmp_k1, m1: tmp_m1, n: nrot, piH: tmp_piH}
                )
                res.append(tmp1)
            return res

        paras_km = list(itertools.product(*[qps_value, m1_value]))
        paras_symbol = [k1, m1, piH]
        characters, paras_values = [], []
        for ii, paras_value in enumerate(paras_km):
            tmp_k1, tmp_m1 = paras_value
            if np.isclose(tmp_k1, 0, atol=symprec):
                idx_fc = 0
                fc = func[idx_fc]
                for tmp_piH in [-1, 1]:
                    res = value_fc(fc, tmp_k1, tmp_m1, tmp_piH, nrot, order)
                    characters.append(np.array(res).astype(np.complex128))
                    paras_values.append([tmp_k1, tmp_m1, tmp_piH])
            else:
                idx_fc = 1
                fc = func[idx_fc]
                tmp_piH = 0
                res = value_fc(fc, tmp_k1, tmp_m1, tmp_piH, nrot, order)
                characters.append(np.array(res).astype(np.complex128))
                paras_values.append([tmp_k1, tmp_m1, tmp_piH])

    elif family == 6:
        qpoint = DictParams["qpoints"]
        nrot = DictParams["nrot"]
        a = DictParams["a"]
        order = DictParams["order"]
        k1, m1, n, piV = symbols("k1 m1 n piV")

        func0 = sympy.Matrix(
            [
                1,
                sympy.exp(1j * k1 * a),
                sympy.exp(1j * m1 * 2 * sympy.pi / n),
                piV,
            ]
        )
        func1 = [
            sympy.Matrix([[1, 0], [0, 1]]),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * k1 * a),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(1j * k1 * a),
                    ],
                ]
            ),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * m1 * 2 * sympy.pi / n),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(-1j * m1 * 2 * sympy.pi / n),
                    ],
                ]
            ),
            sympy.Matrix([[0, 1], [1, 0]]),
        ]
        func = [func0, func1]
        qps_value = [qpoint]
        m1_value = list(range(0, int(nrot / 2) + 1))

        def value_fc(
            fc: list,
            tmp_k1: float,
            tmp_m1: float,
            tmp_piV: int,
            nrot: int,
            order: List[List[int]],
        ) -> list:
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0
                tmp1 = tmp0.evalf(
                    subs={k1: tmp_k1, m1: tmp_m1, n: nrot, piV: tmp_piV}
                )
                res.append(tmp1)
            return res

        paras_km = list(itertools.product(*[qps_value, m1_value]))
        paras_symbol = [k1, m1, piV]
        characters, paras_values = [], []
        for ii, paras_value in enumerate(paras_km):
            tmp_k1, tmp_m1 = paras_value
            if np.isclose(tmp_m1, 0, atol=symprec) or np.isclose(
                tmp_m1, nrot / 2, atol=symprec
            ):
                idx_fc = 0
                fc = func[idx_fc]
                for tmp_piV in [-1, 1]:
                    res = value_fc(fc, tmp_k1, tmp_m1, tmp_piV, nrot, order)
                    characters.append(np.array(res).astype(np.complex128))
                    paras_values.append([tmp_k1, tmp_m1, tmp_piV])
            else:
                idx_fc = 1
                fc = func[idx_fc]
                tmp_piV = 0
                res = value_fc(fc, tmp_k1, tmp_m1, tmp_piV, nrot, order)
                characters.append(np.array(res).astype(np.complex128))
                paras_values.append([tmp_k1, tmp_m1, tmp_piV])

    elif family == 8:
        qpoint = DictParams["qpoints"]
        nrot = DictParams["nrot"]
        a = DictParams["a"]
        order = DictParams["order"]

        k1, m1, n, piV = symbols("k1 m1 n piV")
        func0 = sympy.Matrix(
            [
                1,
                sympy.exp(1j * (k1 * a / 2 + m1 * sympy.pi / n)),
                1,
                piV,
            ]
        )

        func1 = [
            sympy.Matrix([[1, 0], [0, 1]]),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * (k1 * a / 2 + m1 * sympy.pi / n)),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(1j * (k1 * a / 2 - m1 * sympy.pi / n)),
                    ],
                ]
            ),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * m1 * 2 * sympy.pi / n),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(-1j * m1 * 2 * sympy.pi / n),
                    ],
                ]
            ),
            sympy.Matrix([[0, 1], [1, 0]]),
        ]

        func = [func0, func1]
        qps_value = [qpoint]
        m1_value = list(range(0, nrot + 1))

        def value_fc(
            fc: list,
            tmp_k1: float,
            tmp_m1: float,
            tmp_piV: int,
            nrot: int,
            order: List[List[int]],
        ) -> list:
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0
                tmp1 = tmp0.evalf(
                    subs={k1: tmp_k1, m1: tmp_m1, n: nrot, piV: tmp_piV}
                )
                res.append(tmp1)
            return res

        paras_km = list(itertools.product(*[qps_value, m1_value]))
        paras_symbol = [k1, m1, piV]
        characters, paras_values = [], []
        for ii, paras_value in enumerate(paras_km):
            tmp_k1, tmp_m1 = paras_value
            if np.isclose(tmp_m1, 0, atol=symprec) or np.isclose(
                tmp_m1, nrot, atol=symprec
            ):
                idx_fc = 0
                fc = func[idx_fc]
                for tmp_piV in [-1, 1]:
                    res = value_fc(fc, tmp_k1, tmp_m1, tmp_piV, nrot, order)
                    characters.append(np.array(res).astype(np.complex128))
                    paras_values.append([tmp_k1, tmp_m1, tmp_piV])
            else:
                idx_fc = 1
                fc = func[idx_fc]
                tmp_piV = 0
                res = value_fc(fc, tmp_k1, tmp_m1, tmp_piV, nrot, order)
                characters.append(np.array(res).astype(np.complex128))
                paras_values.append([tmp_k1, tmp_m1, tmp_piV])
    elif family == 13:
        qpoint = DictParams["qpoints"]
        nrot = DictParams["nrot"]
        a = DictParams["a"]
        order = DictParams["order"]
        n, k1, m1, f, piU, piV, piH = symbols("n k1 m1 f piU piV piH")

        func0 = sympy.Matrix(
            [
                1,
                sympy.exp(1j * m1 * sympy.pi / n),
                1,
                piU,
                piV,
            ]
        )
        func1 = [
            sympy.Matrix(
                [
                    [1, 0],
                    [0, 1],
                ]
            ),
            sympy.Matrix(
                [
                    [sympy.exp(1j * m1 * sympy.pi / n), 0],
                    [0, sympy.exp(-1j * m1 * sympy.pi / n)],
                ]
            ),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * 2 * m1 * sympy.pi / n),
                        0,
                    ],
                    [
                        0,
                        sympy.exp(-1j * 2 * m1 * sympy.pi / n),
                    ],
                ]
            ),
            sympy.Matrix(
                [
                    [
                        0,
                        piH,
                    ],
                    [
                        piH,
                        0,
                    ],
                ]
            ),
            sympy.Matrix([[0, 1], [1, 0]]),
        ]
        func2 = [
            sympy.Matrix([[1, 0], [0, 1]]),
            sympy.Matrix(
                [
                    [sympy.exp(1j * (m1 * sympy.pi / n + k1 * a / 2)), 0],
                    [0, sympy.exp(1j * (m1 * sympy.pi / n - k1 * a / 2))],
                ]
            ),
            sympy.Matrix(
                [[0, 1], [1, 0]]
            ),  # Suspecting it's an error in the book of line group
            sympy.Matrix([[0, piV], [piV, 0]]),
            sympy.Matrix([[piV, 0], [0, piV]]),
        ]
        func3 = [
            sympy.Matrix([[1, 0], [0, 1]]),
            sympy.Matrix([[-1, 0], [0, 1]]),
            sympy.Matrix([[-1, 0], [0, -1]]),
            sympy.Matrix([[piU, 0], [0, piU]]),
            sympy.Matrix([[0, 1], [1, 0]]),
        ]
        func4 = [
            sympy.Matrix(
                [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
            ),
            sympy.Matrix(
                [
                    [
                        sympy.exp(1j * (m1 * sympy.pi / n + k1 * a / 2)),
                        0,
                        0,
                        0,
                    ],
                    [
                        0,
                        sympy.exp(1j * (-m1 * sympy.pi / n + k1 * a / 2)),
                        0,
                        0,
                    ],
                    [
                        0,
                        0,
                        sympy.exp(1j * (m1 * sympy.pi / n - k1 * a / 2)),
                        0,
                    ],
                    [
                        0,
                        0,
                        0,
                        sympy.exp(-1j * (m1 * sympy.pi / n + k1 * a / 2)),
                    ],
                ]
            ),
            sympy.Matrix(
                [
                    [sympy.exp(1j * 2 * sympy.pi * m1 / n), 0, 0, 0],
                    [0, sympy.exp(-1j * 2 * sympy.pi * m1 / n), 0, 0],
                    [0, 0, sympy.exp(1j * 2 * sympy.pi * m1 / n), 0],
                    [0, 0, 0, sympy.exp(-1j * 2 * sympy.pi * m1 / n)],
                ]
            ),
            sympy.Matrix(
                [[0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]]
            ),
            sympy.Matrix(
                [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]
            ),
        ]

        func = [func0, func1, func2, func3, func4]
        qps_value = [qpoint]

        if np.isclose(np.abs(qpoint), np.pi / a, atol=symprec):
            m1_value = list(range(0, np.floor(nrot / 2 + 1).astype(np.int32)))
        else:
            m1_value = list(range(0, nrot + 1))

        def value_fc(
            fc: list,
            tmp_k1: float,
            tmp_m1: float,
            tmp_piU: int,
            tmp_piV: int,
            tmp_piH: int,
            nrot: int,
            a: float,
            order: List[List[int]],
        ) -> list:
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = tmp0 * fc[tmp]
                tmp1 = tmp0.xreplace(
                    {
                        k1: tmp_k1,
                        m1: tmp_m1,
                        n: nrot,
                        f: a / 2,
                        piU: tmp_piU,
                        piV: tmp_piV,
                        piH: tmp_piH,
                    }
                )
                res.append(tmp1)
            return res

        paras_km = list(itertools.product(*[qps_value, m1_value]))
        paras_symbol = [k1, m1, piU, piV, piH]
        characters, paras_values = [], []
        for ii, paras_value in enumerate(paras_km):
            tmp_k1, tmp_m1 = paras_value
            if np.isclose(tmp_k1, 0, atol=symprec):
                if np.isclose(np.abs(tmp_m1), 0) or np.isclose(
                    np.abs(tmp_m1), nrot
                ):
                    idx_fc = 0
                    fc = func[idx_fc]
                    tmp_piH = 0
                    for tmp_piU, tmp_piV in itertools.product(
                        [-1, 1], [-1, 1]
                    ):
                        res = value_fc(
                            fc,
                            tmp_k1,
                            tmp_m1,
                            tmp_piU,
                            tmp_piV,
                            tmp_piH,
                            nrot,
                            a,
                            order,
                        )
                        characters.append(np.array(res).astype(np.complex128))
                        paras_values.append(
                            [tmp_k1, tmp_m1, tmp_piU, tmp_piV, tmp_piH]
                        )
                else:
                    idx_fc = 1
                    fc = func[idx_fc]
                    tmp_piU, tmp_piV = 0, 0
                    for tmp_piH in [-1, 1]:
                        res = value_fc(
                            fc,
                            tmp_k1,
                            tmp_m1,
                            tmp_piU,
                            tmp_piV,
                            tmp_piH,
                            nrot,
                            a,
                            order,
                        )
                        characters.append(np.array(res).astype(np.complex128))
                        paras_values.append(
                            [tmp_k1, tmp_m1, tmp_piU, tmp_piV, tmp_piH]
                        )
            elif np.isclose(np.abs(tmp_k1), np.pi / a, atol=symprec):
                if np.isclose(np.abs(tmp_m1), nrot / 2, atol=symprec):
                    idx_fc = 3
                    fc = func[idx_fc]
                    tmp_piV, tmp_piH = 0, 0
                    for tmp_piU in [-1, 1]:
                        res = value_fc(
                            fc,
                            tmp_k1,
                            tmp_m1,
                            tmp_piU,
                            tmp_piV,
                            tmp_piH,
                            nrot,
                            a,
                            order,
                        )
                        characters.append(np.array(res).astype(np.complex128))
                        paras_values.append(
                            [tmp_k1, tmp_m1, tmp_piU, tmp_piV, tmp_piH]
                        )

                elif np.isclose(np.abs(tmp_m1), 0, atol=symprec):
                    idx_fc = 2
                    fc = func[idx_fc]
                    tmp_piU, tmp_piH = 0, 0
                    for tmp_piV in [-1, 1]:
                        res = value_fc(
                            fc,
                            tmp_k1,
                            tmp_m1,
                            tmp_piU,
                            tmp_piV,
                            tmp_piH,
                            nrot,
                            a,
                            order,
                        )
                        characters.append(np.array(res).astype(np.complex128))
                        paras_values.append(
                            [tmp_k1, tmp_m1, tmp_piU, tmp_piV, tmp_piH]
                        )
                else:
                    idx_fc = 4
                    fc = func[idx_fc]
                    tmp_piU, tmp_piV, tmp_piH = 0, 0, 0
                    res = value_fc(
                        fc,
                        tmp_k1,
                        tmp_m1,
                        tmp_piU,
                        tmp_piV,
                        tmp_piH,
                        nrot,
                        a,
                        order,
                    )
                    characters.append(np.array(res).astype(np.complex128))
                    paras_values.append(
                        [tmp_k1, tmp_m1, tmp_piU, tmp_piV, tmp_piH]
                    )

            elif 0 < np.abs(tmp_k1) < np.pi / a:
                if np.isclose(np.abs(tmp_m1), 0, atol=symprec) or np.isclose(
                    np.abs(tmp_m1), nrot, atol=symprec
                ):
                    idx_fc = 2
                    fc = func[idx_fc]
                    tmp_piU, tmp_piH = 0, 0
                    for tmp_piV in [-1, 1]:
                        res = value_fc(
                            fc,
                            tmp_k1,
                            tmp_m1,
                            tmp_piU,
                            tmp_piV,
                            tmp_piH,
                            nrot,
                            a,
                            order,
                        )
                        characters.append(np.array(res).astype(np.complex128))
                        paras_values.append(
                            [tmp_k1, tmp_m1, tmp_piU, tmp_piV, tmp_piH]
                        )

                else:
                    idx_fc = 4
                    fc = func[idx_fc]
                    tmp_piU, tmp_piV, tmp_piH = 0, 0, 0
                    res = value_fc(
                        fc,
                        tmp_k1,
                        tmp_m1,
                        tmp_piU,
                        tmp_piV,
                        tmp_piH,
                        nrot,
                        a,
                        order,
                    )
                    characters.append(np.array(res).astype(np.complex128))
                    paras_values.append(
                        [tmp_k1, tmp_m1, tmp_piU, tmp_piV, tmp_piH]
                    )
            else:
                logging.ERROR("Wrong value for k1")
    else:
        raise NotImplementedError("Family %d is not supported yet" % family)
    return characters, paras_values, paras_symbol
