# 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 Tuple

import numpy as np
import sympy
from sympy import symbols


def frac_range(
    start: float,
    end: float,
    left: bool = True,
    right: bool = True,
    symprec: float = 0.01,
) -> list:
    """return the integer within the specified range

    Args:
        start: left boundary
        end: right boundary
        left: False mean delete the left boundary element if it is an integer
        right: False mean delete the right boundary element if it is an integer
        symprec: system precise

    Returns:

    """
    close = list(
        range(
            np.ceil(start).astype(np.int32), np.floor(end).astype(np.int32) + 1
        )
    )
    if left == False:
        if close[0] - start < symprec:
            close.pop(0)  # delete the left boundary
    if right == False:
        if close[-1] - end < symprec:
            close.pop()  # delete the right boundary
    return close


def line_group_sympy(
    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]
        n_value = [nrot]
        m1_value = list(range(int(-nrot / 2) + 1, int(nrot / 2) + 1))

        if np.isclose(qpoint, 0, atol=symprec) or np.isclose(
            qpoint, np.pi / a, atol=symprec
        ):
            fc = func[0]
            paras_symbol = [k1, m1, n, piH]
        else:
            fc = func[1]
            paras_symbol = [k1, m1, n]
        paras_values = list(itertools.product(*[qps_value, m1_value, n_value]))
        characters = []
        for ii, paras_value in enumerate(paras_values):
            tmp_k1, tmp_m1, tmp_n = paras_value
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = tmp0 * fc[tmp]
                if len(paras_symbol) == 4:
                    tmp1 = tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: -1}
                    ) + tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: 1}
                    )
                    res.append(tmp1)
                else:
                    tmp1 = tmp0.evalf(subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n})
                    res.append(tmp1)
            characters.append(np.array(res).astype(np.complex128))
    elif family == 3:
        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 * 2 * sympy.pi / n),
                piH,
            ]
        )
        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]
        n_value = [nrot]
        m1_value = list(range(int(-nrot / 2) + 1, int(nrot / 2) + 1))

        if np.isclose(qpoint, 0, atol=symprec) or np.isclose(
            qpoint, np.pi / a, atol=symprec
        ):
            fc = func[0]
            paras_symbol = [k1, m1, n, piH]
        else:
            fc = func[1]
            paras_symbol = [k1, m1, n]
        paras_values = list(itertools.product(*[qps_value, m1_value, n_value]))
        characters = []
        for ii, paras_value in enumerate(paras_values):
            tmp_k1, tmp_m1, tmp_n = paras_value
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = tmp0 * fc[tmp]
                if len(paras_symbol) == 4:
                    tmp1 = tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: -1}
                    ) + tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: 1}
                    )
                    res.append(tmp1)
                else:
                    tmp1 = tmp0.evalf(subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n})
                    res.append(tmp1)
            characters.append(np.array(res).astype(np.complex128))
    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]
        n_value = [nrot]

        m1_value = list(range(-nrot + 1, nrot + 1))
        if np.isclose(qpoint, 0, atol=symprec):
            idx_fc = 0
            fc = func[idx_fc]
            paras_symbol = [k1, m1, n, piH]
        else:
            idx_fc = 1
            fc = func[idx_fc]
            paras_symbol = [k1, m1, n]
        paras_values = list(itertools.product(*[qps_value, m1_value, n_value]))
        characters = []
        for ii, paras_value in enumerate(paras_values):
            tmp_k1, tmp_m1, tmp_n = paras_value
            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0
                if idx_fc == 0:
                    tmp1 = tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: -1}
                    ) + tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: 1}
                    )
                    res.append(tmp1)
                else:
                    tmp1 = tmp0.evalf(subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n})
                    res.append(tmp1)
            characters.append(np.array(res).astype(np.complex128))
    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]
        n_value = [nrot]
        m1_value = list(range(0, int(nrot / 2) + 1))

        paras_values = list(itertools.product(*[qps_value, m1_value, n_value]))
        characters = []
        for ii, paras_value in enumerate(paras_values):
            tmp_k1, tmp_m1, tmp_n = paras_value

            if np.isclose(tmp_m1, 0, atol=symprec) or np.isclose(
                tmp_m1, tmp_n / 2, atol=symprec
            ):
                fc = func[0]
                paras_symbol = [k1, m1, n, piV]
            else:
                fc = func[1]
                paras_symbol = [k1, m1, n]

            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0
                if len(paras_symbol) == 4:
                    tmp1 = tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piV: -1}
                    ) + tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piV: 1}
                    )
                    res.append(tmp1)
                else:
                    tmp1 = tmp0.evalf(subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n})
                    res.append(tmp1)
            characters.append(np.array(res).astype(np.complex128))
    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]
        n_value = [nrot]

        m1_value = list(range(0, nrot + 1))
        paras_values = list(itertools.product(*[qps_value, m1_value, n_value]))
        characters = []
        for ii, paras_value in enumerate(paras_values):
            tmp_k1, tmp_m1, tmp_n = paras_value

            if np.isclose(tmp_m1, 0, atol=symprec) or np.isclose(
                tmp_m1, nrot, atol=symprec
            ):
                fc = func[0]
                paras_symbol = [k1, m1, n, piV]
            else:
                fc = func[1]
                paras_symbol = [k1, m1, n]

            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0
                if len(paras_symbol) == 4:
                    tmp1 = tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piV: -1}
                    ) + tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piV: 1}
                    )
                    res.append(tmp1)
                else:
                    tmp1 = tmp0.evalf(subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n})
                    res.append(tmp1)
            characters.append(np.array(res).astype(np.complex128))
    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]]),
            sympy.Matrix([[1, 0], [0, 1]]),
            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]
        n_value = [nrot]

        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))

        f_value = [a / 2]
        paras_values = list(itertools.product(*[qps_value, m1_value, n_value]))
        characters = []
        for ii, paras_value in enumerate(paras_values):
            tmp_k1, tmp_m1, tmp_n = paras_value
            # tmp_k1, tmp_m1, tmp_n  = 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), tmp_n
                ):
                    idx_fc = 0
                    fc = func[idx_fc]
                    paras_symbol = [k1, m1, piU, piV]
                elif 0 < np.abs(tmp_m1) < tmp_n:
                    idx_fc = 1
                    fc = func[idx_fc]
                    paras_symbol = [k1, m1, piH]
            elif np.isclose(np.abs(tmp_k1), np.pi / a, atol=symprec):
                if np.isclose(np.abs(tmp_m1), tmp_n / 2, atol=symprec):
                    idx_fc = 3
                    fc = func[idx_fc]
                    paras_symbol = [piU]
                elif np.isclose(np.abs(tmp_m1), 0, atol=symprec):
                    idx_fc = 2
                    fc = func[idx_fc]
                    paras_symbol = [k1, m1, piV]
                elif 0 < np.abs(tmp_m1) < tmp_n / 2:
                    idx_fc = 4
                    fc = func[idx_fc]
                    paras_symbol = [k1, m1, n]
            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), tmp_n, atol=symprec
                ):
                    idx_fc = 2
                    fc = func[idx_fc]
                    paras_symbol = [k1, m1, piV]

                elif 0 < np.abs(tmp_m1) < tmp_n:
                    idx_fc = 4
                    fc = func[idx_fc]
                    paras_symbol = [k1, m1]
            else:
                logging.ERROR("Wrong value for k1")

            res = []
            for tmp_order in order:
                for jj, tmp in enumerate(tmp_order):
                    if jj == 0:
                        tmp0 = fc[tmp]
                    else:
                        tmp0 = fc[tmp] * tmp0

                if idx_fc == 0:
                    tmp1 = (
                        tmp0.evalf(
                            subs={
                                k1: tmp_k1,
                                m1: tmp_m1,
                                n: tmp_n,
                                piV: -1,
                                piU: 1,
                            }
                        )
                        + tmp0.evalf(
                            subs={
                                k1: tmp_k1,
                                m1: tmp_m1,
                                n: tmp_n,
                                piV: 1,
                                piU: 1,
                            }
                        )
                        + tmp0.evalf(
                            subs={
                                k1: tmp_k1,
                                m1: tmp_m1,
                                n: tmp_n,
                                piV: -1,
                                piU: -1,
                            }
                        )
                        + tmp0.evalf(
                            subs={
                                k1: tmp_k1,
                                m1: tmp_m1,
                                n: tmp_n,
                                piV: 1,
                                piU: -1,
                            }
                        )
                    )
                    res.append(tmp1)
                elif idx_fc == 1:
                    tmp1 = tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: -1}
                    ) + tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piH: 1}
                    )
                    res.append(tmp1)
                elif idx_fc == 2:
                    tmp1 = tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piV: -1}
                    ) + tmp0.evalf(
                        subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n, piV: 1}
                    )
                    res.append(tmp1)
                elif idx_fc == 3:
                    tmp1 = tmp0.evalf(subs={piU: -1}) + tmp0.evalf(
                        subs={piU: 1}
                    )
                    res.append(tmp1)
                elif idx_fc == 4:
                    tmp1 = tmp0.evalf(subs={k1: tmp_k1, m1: tmp_m1, n: tmp_n})
                    res.append(tmp1)
            characters.append(np.array(res).astype(np.complex128))
    else:
        raise NotImplementedError("Family %d is not supported yet" % family)
    return characters, paras_values, paras_symbol
