"""Set of useful utility functions."""

from collections import OrderedDict
from collections import defaultdict
from scipy import sparse
from scipy.sparse.linalg import spsolve
import random
import numpy as np
import math
import xmltodict
import re
import requests


def xml_to_dict(fname):
    """Parse XML file."""
    with open(fname, "r") as f:
        data = xmltodict.parse(f.read())
    return data


def get_factors(x, start=2):
    """Get factors of a number."""
    facts = []
    for i in range(start, x + 1):
        if x % i == 0:
            facts.append(i)
    return facts


def get_counts(array=["W", "W", "Mo", "Mo", "S", "S"]):
    """
    Get number of unique elements and their counts.

    Uses OrderedDict.

    Args:
         array of elements

    Returns:
         ordereddict, e.g.OrderedDict([('W', 2), ('Mo', 2), ('S', 2)])
    """
    uniqe_els = []
    for i in array:
        if i not in uniqe_els:
            uniqe_els.append(i)
    info = OrderedDict()
    for i in uniqe_els:
        info.setdefault(i, 0)
    for i in array:
        info[i] += 1
    return info


def gcd(a, b):
    """Calculate the Greatest Common Divisor of a and b.

    Unless b==0, the result will have the same sign as b (so that when
    b is divided by it, the result comes out positive).
    """
    while b:
        a, b = b, a % b
    return a


def ext_gcd(a, b):
    """GCD module from ase."""
    if b == 0:
        return 1, 0
    elif a % b == 0:
        return 0, 1
    else:
        x, y = ext_gcd(b, a % b)
        return y, x - y * (a // b)


def rand_select(x=[]):
    """Select randomly with index info."""
    info = {}
    for i, ii in enumerate(x):
        info.setdefault(ii, []).append(i)
    selected = {}
    for i, j in info.items():
        chosen = random.choice(j)
        selected.setdefault(i, chosen)
    return selected


def rec_dict():
    """Make a recursion dictionary."""
    return defaultdict(rec_dict)


def random_colors(number_of_colors=110):
    """Generate random colors for atom coloring."""
    colors = [
        "#" + "".join([random.choice("0123456789ABCDEF") for j in range(6)])
        for i in range(number_of_colors)
    ]
    color_dict = {}
    for i, ii in enumerate(colors):
        color_dict[i] = ii
    return color_dict


def get_angle(
    a=np.array([1, 2, 3]), b=np.array([4, 5, 6]), c=np.array([7, 8, 9])
):
    """Get angle between three vectors."""
    # theta = argcos(x.y/(|x||y|))
    cos = np.dot((a - b), (c - b)) / (
        np.linalg.norm((a - b)) * np.linalg.norm((c - b))
    )
    if cos <= -1.0:
        cos = cos + 0.000001
    if cos >= 1.0:
        cos = cos - 0.000001
    angle = math.degrees(math.acos(cos))
    return angle


def recast_array_on_uniq_array_elements(
    uniq=["Si", "Al", "O"],
    arr=["Si", "Si", "Al", "Al", "Si", "O", "O", "O", "O"],
):
    """Recast array on uniq array elements."""
    info = {}
    for i, ii in enumerate(uniq):
        for j, jj in enumerate(arr):
            if ii == jj:
                info.setdefault(ii, []).append(j)
    return info


def lorentzian(x, y, x0, gamma):
    """Get Lorentzian of a function."""
    return (y / math.pi) * (
        (0.5 * gamma) / ((x - x0) ** 2 + (0.5 * gamma) ** 2)
    )


# color_dict=random_colors()
def stringdict_to_xml(d={}, enforce_string=False):
    """Convert string dictionary to XML."""
    line = ""
    for i, j in d.items():
        if enforce_string:
            line += "<" + str(i) + ">'" + str(j) + "'</" + str(i) + ">"
        else:
            line += "<" + str(i) + ">" + str(j) + "</" + str(i) + ">"
    return line


def array_to_string(arr=[]):
    """Convert 1D arry to string."""
    return ",".join(map(str, arr))


def chunks(lst, n):
    """Split successive n-sized chunks from list."""
    x = []
    for i in range(0, len(lst), n):
        x.append(lst[i : i + n])
    return x


def check_match(a, b, tol=1e-8):
    """Check if a and b are the same, taking into account PBCs."""
    if abs(a[0] - b[0]) < tol or abs(abs(a[0] - b[0]) - 1) < tol:
        if abs(a[1] - b[1]) < tol or abs(abs(a[1] - b[1]) - 1) < tol:
            if abs(a[2] - b[2]) < tol or abs(abs(a[2] - b[2]) - 1) < tol:
                return True
    return False


def update_dict(main={}, extra={}):
    """Return update dictionary."""
    # Helper function for dict.update method
    tmp = main.copy()
    for i, j in extra.items():
        tmp[i] = j
    return tmp


def get_new_coord_for_xyz_sym(frac_coord=[], xyz_string=""):
    """Obtain new coord from xyz string."""
    affine_matrix = parse_xyz_string(xyz_string)
    coord = operate_affine(frac_coord, affine_matrix)
    coord = np.array([i - math.floor(i) for i in coord])
    return coord


def check_duplicate_coords(coords=[], coord=[]):
    """Check if a coordinate exists."""
    positive = False
    for i in coords:
        tmp = check_match(i, coord)
        if tmp:
            positive = True
    return positive


def parse_xyz_string(xyz_string):
    """
    Convert xyz info to translation and rotation vectors.

    Adapted from pymatgen.
    Args:
        xyz_string: string of the form 'x, y, z', '-x, -y, z',
            '-2y+1/2, 3x+1/2, z-y+1/2', etc.
    Returns:
        translation and rotation vectors.
    """
    rot_matrix = np.zeros((3, 3))
    trans = np.zeros(3)
    toks = xyz_string.strip().replace(" ", "").lower().split(",")
    re_rot = re.compile(r"([+-]?)([\d\.]*)/?([\d\.]*)([x-z])")
    re_trans = re.compile(r"([+-]?)([\d\.]+)/?([\d\.]*)(?![x-z])")
    for i, tok in enumerate(toks):
        # build the rotation matrix
        for m in re_rot.finditer(tok):
            factor = -1 if m.group(1) == "-" else 1
            if m.group(2) != "":
                factor *= (
                    float(m.group(2)) / float(m.group(3))
                    if m.group(3) != ""
                    else float(m.group(2))
                )
            j = ord(m.group(4)) - 120
            rot_matrix[i, j] = factor
        # build the translation vector
        for m in re_trans.finditer(tok):
            factor = -1 if m.group(1) == "-" else 1
            num = (
                float(m.group(2)) / float(m.group(3))
                if m.group(3) != ""
                else float(m.group(2))
            )
            trans[i] = num * factor
    affine_matrix = np.eye(4)
    affine_matrix[0:3][:, 0:3] = rot_matrix
    affine_matrix[0:3][:, 3] = trans
    return affine_matrix


def operate_affine(cart_coord=[], affine_matrix=[]):
    """Operate affine method."""
    affine_point = np.array([cart_coord[0], cart_coord[1], cart_coord[2], 1])
    return np.dot(np.array(affine_matrix), affine_point)[0:3]


def gaussian(x, sigma):
    """Get Gaussian profile."""
    return np.exp(-(x**2) / (2 * sigma**2))


def lorentzian2(x, gamma):
    """Get Lorentziann profile."""
    return (
        gamma / 2 / (np.pi * (x**2 + (gamma / 2) ** 2)) / (2 / (np.pi * gamma))
    )


def digitize_array(values=[], max_len=10):
    """Digitze an array."""
    has_float = False in [float(i).is_integer() for i in values]
    if has_float:
        arr = np.array([float(i) for i in values])
        max_val = max(arr)
        min_val = min(arr)
        bins = np.arange(1, max_len + 1) * (max_val - min_val) / 10
        values = np.digitize(arr, bins)
    return values


def bond_angle(
    dist1,
    dist2,
    bondx1,
    bondx2,
    bondy1,
    bondy2,
    bondz1,
    bondz2,
):
    """Get an angle."""
    nm = dist1 * dist2
    rrx = bondx1 * bondx2
    rry = bondy1 * bondy2
    rrz = bondz1 * bondz2
    cos = (rrx + rry + rrz) / (nm)
    if cos <= -1.0:
        cos = cos + 0.000001
    if cos >= 1.0:
        cos = cos - 0.000001
    deg = math.degrees(math.acos(cos))
    return deg


def check_url_exists(
    url="https://www.ctcms.nist.gov/~knc6/static/JARVIS-DFT/JVASP-77580.xml",
):
    """Check if a url exists."""
    request = requests.get(url)
    if request.status_code == 200:
        return True
    else:
        return False


def volumetric_grid_reshape(data=[], final_grid=[50, 50, 50]):
    """Reshape volumetric data."""
    import torch

    data = torch.tensor(data).unsqueeze(0).unsqueeze(0)
    new_data = (
        torch.nn.functional.interpolate(
            data,
            size=final_grid,
            scale_factor=None,
            mode="trilinear",
            align_corners=True,
            recompute_scale_factor=None,
        )
        .squeeze()
        .squeeze()
    )
    return new_data.numpy()


def cos_formula(a, b, c):
    """Get angle between three edges for oblique triangles."""
    res = (a**2 + b**2 - c**2) / (2 * a * b)
    res = -1.0 if res < -1.0 else res
    res = 1.0 if res > 1.0 else res
    return np.arccos(res)


def baseline_als(y, lam, p, niter=10):
    """
    Adaptive Least Squares fitting for baseline correction

    Parameters:
    y: array_like
        Input signal
    lam: float
        Lambda (smoothness)
    p: float
        Asymmetry
    niter: int, optional
        Number of iterations

    Returns:
    array_like
        The estimated baseline
    """
    L = len(y)
    D = sparse.diags([1, -2, 1], [0, -1, -2], shape=(L, L - 2))
    w = np.ones(L)
    for i in range(niter):
        W = sparse.spdiags(w, 0, L, L)
        Z = W + lam * D.dot(D.transpose())
        z = spsolve(Z, w * y)
        w = p * (y > z) + (1 - p) * (y < z)
    return z


def recast_array(
    x_original=[], y_original=[], x_new=np.arange(0, 90, 1), tol=0.1
):
    x_original = np.array(x_original)
    # Initialize the new y array with NaNs or a default value
    y_new = np.full_like(x_new, 0, dtype=np.float64)

    # Fill the corresponding bins
    for x_val, y_val in zip(x_original, y_original):
        closest_index = np.abs(
            x_new - x_val
        ).argmin()  # Find the closest x_new index
        y_new[closest_index] = y_val
    # y_new[y_new<tol]=0
    return x_new, y_new


# def is_xml_valid(xsd="jarvisdft.xsd", xml="JVASP-1002.xml"):
#   """Check if XML is valid."""
#   xml_file = etree.parse(xml)
#   xml_validator = etree.XMLSchema(file=xsd)
#   is_valid = xml_validator.validate(xml_file)
#   return is_valid
