import copy
import os
import traceback
from collections import defaultdict
from pathlib import Path

import ase
import ase.neighborlist
import ase.visualize
import numpy as np

from .log import logger


class Framework:
    def __init__(self, atoms, bonds, bond_types, info, wrap=True):
        # Save data to attributes.
        self.atoms = atoms.copy()
        self.bonds = bonds.copy()
        self.bond_types = copy.deepcopy(bond_types)
        self.info = copy.deepcopy(info)

        if wrap:
            self.wrap()

    def wrap(self):
        # Cleanup errors.
        s = self.atoms.get_scaled_positions()
        while (s >= 1.0).any() and (s < 0.0).any():
            s = np.where(
                s < 0.0,
                s + np.ones_like(s),
                s,
            )
            s = np.where(
                s >= 1.0,
                s - np.ones_like(s),
                s,
            )

        s[np.abs(s - 1) < 1e-3] = 0.9999
        s[np.abs(s) < 1e-3] = 0.0001

        self.atoms.set_scaled_positions(s)

    def write_cif(self, filename):
        """
        Write framework in cif format.
        """
        path = Path(filename).resolve()
        if not path.parent.exists():
            logger.error(f"Path {path} does not exist.")

        # Add suffix if not exists.
        if path.suffix != ".cif":
            path = path.with_suffix(".cif")

        try:
            self._write_cif(path)
        except Exception as e:
            logger.error(e)
            logger.error(
                "CIF writing fails with error: %s",
                traceback.format_exc(),
            )
            # Remove invalid cif file.
            logger.error("Remove invalid CIF: %s", path)
            os.remove(str(path))

    def _write_cif(self, path):
        stem = path.stem.replace(" ", "_")

        with path.open("w") as f:
            # Write information comments.
            _topology = self.info["topology"]
            info_str = "# Generated by pormake.\n"
            info_str += "# {}\n".format(_topology)
            info_str += "# Building blocks:\n"
            for i in _topology.node_indices:
                t = _topology.node_types[i]
                bb = self.info["located_bbs"][i]
                info_str += "#     Node {}, Type {}, {}\n".format(i, t, bb)
            for i in _topology.edge_indices:
                t = tuple(_topology.edge_types[i])
                bb = self.info["located_bbs"][i]
                info_str += "#     Edge {}, Type {}, {}\n".format(i, t, bb)
            f.write(info_str)
            f.write(
                "# Relax obj value: {:.3f}\n".format(self.info["relax_obj"])
            )
            f.write("# Max RMSD: {:.3f}\n".format(self.info["max_rmsd"]))
            f.write("# Mean RMSD: {:.3f}\n".format(self.info["mean_rmsd"]))

            f.write("data_{}\n".format(stem))

            f.write("_symmetry_space_group_name_H-M    P1\n")
            f.write("_symmetry_Int_Tables_number       1\n")
            f.write("_symmetry_cell_setting            triclinic\n")

            f.write("loop_\n")
            f.write("_symmetry_equiv_pos_as_xyz\n")
            f.write("'x, y, z'\n")

            (
                a,
                b,
                c,
                alpha,
                beta,
                gamma,
            ) = self.atoms.get_cell_lengths_and_angles()

            f.write("_cell_length_a     {:.3f}\n".format(a))
            f.write("_cell_length_b     {:.3f}\n".format(b))
            f.write("_cell_length_c     {:.3f}\n".format(c))
            f.write("_cell_angle_alpha  {:.3f}\n".format(alpha))
            f.write("_cell_angle_beta   {:.3f}\n".format(beta))
            f.write("_cell_angle_gamma  {:.3f}\n".format(gamma))

            f.write("loop_\n")
            f.write("_atom_site_label\n")
            f.write("_atom_site_type_symbol\n")
            f.write("_atom_site_fract_x\n")
            f.write("_atom_site_fract_y\n")
            f.write("_atom_site_fract_z\n")
            f.write("_atom_type_partial_charge\n")

            symbols = self.atoms.symbols
            charges = self.atoms.get_initial_charges()
            frac_coords = self.atoms.get_scaled_positions()
            for i, (sym, pos, charge) in enumerate(zip(symbols, frac_coords, charges)):
                label = "{}{}".format(sym, i)
                f.write(
                    "{} {} {:.5f} {:.5f} {:.5f} {:.5f}\n".format(label, sym, *pos, charge)
                )

            f.write("loop_\n")
            f.write("_geom_bond_atom_site_label_1\n")
            f.write("_geom_bond_atom_site_label_2\n")
            f.write("_geom_bond_distance\n")
            f.write("_geom_bond_site_symmetry_2\n")
            f.write("_ccdc_geom_bond_type\n")  # ?????????

            # Make bond type dict.
            bond_type_dict = {}
            for (i, j), t in zip(self.bonds, self.bond_types):
                bond_type_dict[(i, j)] = t

            # Get images and distances.
            I, J, S, D = ase.neighborlist.neighbor_list(
                "ijSd", self.atoms, cutoff=6.0
            )
            image_dict = {}
            distance_dict = defaultdict(lambda: 1e30)
            origin = np.array([5, 5, 5])
            for i, j, s, d in zip(I, J, S, D):
                # Take minimum image only.
                if d < distance_dict[(i, j)]:
                    image_dict[(i, j)] = origin + s
                    distance_dict[(i, j)] = d

            for bond in self.bonds:
                i, j = bond

                sym = self.atoms.symbols[i]
                label_i = "{}{}".format(sym, i)

                sym = self.atoms.symbols[j]
                label_j = "{}{}".format(sym, j)

                distance = distance_dict[(i, j)]
                bond_type = bond_type_dict[(i, j)]

                image = image_dict[(i, j)]

                if (image == origin).all():
                    f.write(
                        "{} {} {:.3f} . {}\n".format(
                            label_i, label_j, distance, bond_type
                        )
                    )
                else:
                    f.write(
                        "{} {} {:.3f} 1_{}{}{} {}\n".format(
                            label_i, label_j, distance, *image, bond_type
                        )
                    )

    def view(self, *args, **kwargs):
        ase.visualize.view(self.atoms, *args, **kwargs)
