import json
import logging
import os
from time import time
from xml.etree import ElementTree as ET

from gemmi import cif

from pisa_utils.dictionaries import get_assembly_dict, get_bond_dict, get_molecules_dict

LOGGER = logging.getLogger(__name__)


class AnalysePisa:
    def __init__(self, pdb_id, assembly_id, input_updated_cif):
        """Parses XML files generated by pisa, and serializing to JSON format.

        Also adds UniProt numbering using updated mmCIF file if provided.

        Args:
            pdb_id (str): PDB ID.
            assembly_id (str): Assembly ID.
            input_updated_cif (str): Path to updated mmCIF file with Uniprot Mapping.
        """
        self.pdb_id = pdb_id
        self.assembly_code = assembly_id
        self.input_updated_cif = input_updated_cif

    def interfaces_xml_to_json(self, assembly_xml, interfaces_xml, output_json):
        """Extracts information from PISA assembly and interface XML files and
        serializes to JSON.

        JSON is written as dict of lists to make it compact for the web app.

        Args:
            assembly_xml (str): Path to assembly XML file.
            interfaces_xml (str): Path to interfaces XML file.
            output_json (str): Path to the output interface JSON file.
        """

        start = time()

        logging.info("Parsing XML files {} and {}".format(assembly_xml, interfaces_xml))

        asroot = ET.parse(assembly_xml).getroot()
        root = ET.parse(interfaces_xml).getroot()

        assemblies = asroot.iter("asu_complex")
        assem_result = get_assembly_dict(assemblies)

        # Create interfaces dictionaries
        result = {}
        result["status"] = root.find("status").text
        result["num_interfaces"] = root.find("n_interfaces").text
        logging.info(f"Total no. of interfaces: {result['num_interfaces']}")

        valid_interface_count = 0

        interfaces = root.iter("interface")

        if self.input_updated_cif and os.path.exists(self.input_updated_cif):
            doc = cif.read(self.input_updated_cif)
            updated_cif_block = doc.sole_block()
        else:
            # NOTE Should this be a warning/error instead?
            logging.info(
                f"Updated CIF [{self.input_updated_cif}] is not provided or is invalid"
            )
            updated_cif_block = None

        for interface in interfaces:
            # Interface General information
            interface_id = interface.find("id").text
            interface_area = round(float(interface.find("int_area").text), 2)
            interface_solvation_energy = round(
                float(interface.find("int_solv_en").text), 2
            )
            interface_stabilization_energy = round(
                float(interface.find("stab_en").text), 2
            )
            p_value = round(float(interface.find("pvalue").text), 3)

            # No. of bonds counted
            n_h_bonds = int(interface.find("h-bonds/n_bonds").text)
            n_ss_bonds = int(interface.find("ss-bonds/n_bonds").text)
            n_covalent_bonds = int(interface.find("cov-bonds/n_bonds").text)
            n_salt_bridges = int(interface.find("salt-bridges/n_bonds").text)
            other_contacts = int(interface.find("other-bonds/n_bonds").text)

            molecules = interface.iter("molecule")
            molecules_dicts, interface_residues_count, invalid = get_molecules_dict(
                molecules
            )
            if invalid:
                logging.info(f"Invalid interface: {interface_id}. Skipping...")
                continue

            if (
                n_h_bonds == 0
                and n_ss_bonds == 0
                and n_covalent_bonds == 0
                and n_salt_bridges == 0
                and other_contacts == 0
            ):
                logging.info(
                    "Bonds: (n_h_bonds, n_ss_bonds, n_covalent_bonds, n_salt_bridges)"
                )
                logging.info(
                    f"Bonds: ({n_h_bonds}, {n_ss_bonds}, {n_covalent_bonds}, "
                    f"{n_salt_bridges})"
                )
                logging.info(f"Invalid interface: {interface_id}. Skipping...")
                continue

            valid_interface_count += 1

            # Reading bonds

            hbonds = interface.find("h-bonds")
            sbridges = interface.find("salt-bridges")
            covbonds = interface.find("cov-bonds")
            ssbonds = interface.find("ss-bonds")
            othbonds = interface.find("other-bonds")

            # Writing bonds dictionaries

            # for hbond in hbonds:
            hbond_dict = get_bond_dict(hbonds, "H-bond", updated_cif_block)

            # for sbridge in sbridges:
            sbridge_dict = get_bond_dict(sbridges, "salt-bridges", updated_cif_block)
            # for covbond in covbonds:
            covbond_dict = get_bond_dict(covbonds, "cov-bonds", updated_cif_block)
            # for ssbond in ssbonds:
            ssbond_dict = get_bond_dict(ssbonds, "ss-bonds", updated_cif_block)
            # for othbond in othbonds:
            othbond_dict = get_bond_dict(othbonds, "other-bond", updated_cif_block)

            interface_dict = {
                "interface_id": interface_id,
                "interface_area": interface_area,
                "solvation_energy": interface_solvation_energy,
                "stabilization_energy": interface_stabilization_energy,
                "p_value": p_value,
                "number_interface_residues": interface_residues_count,
                "number_hydrogen_bonds": n_h_bonds,
                "number_covalent_bonds": n_covalent_bonds,
                "number_disulfide_bonds": n_ss_bonds,
                "number_salt_bridges": n_salt_bridges,
                "number_other_bonds": other_contacts,
                "hydrogen_bonds": hbond_dict,
                "salt_bridges": sbridge_dict,
                "disulfide_bonds": ssbond_dict,
                "covalent_bonds": covbond_dict,
                "other_bonds": othbond_dict,
                "molecules": molecules_dicts,
            }

            # Append all dictionaries in 'result'
            result.setdefault("id", []).append(interface_id)
            result.setdefault("int_area", []).append(interface_area)
            result.setdefault("interface_dicts", []).append(interface_dict)

        result["valid_interface_count"] = valid_interface_count
        logging.info(f"valid interface count: {valid_interface_count}")

        interface_count = valid_interface_count

        interface_dicts = result.get("interface_dicts", [])
        assem_dict = {
            "mmsize": assem_result.get("assembly_mmsize"),
            "dissociation_energy": assem_result.get("assembly_diss_energy"),
            "accessible_surface_area": assem_result.get("assembly_asa"),
            "buried_surface_area": assem_result.get("assembly_bsa"),
            "entropy": assem_result.get("assembly_entropy"),
            "dissociation_area": assem_result.get("assembly_diss_area"),
            "solvation_energy_gain": assem_result.get("assembly_int_energy"),
            "formula": assem_result.get("assembly_formula"),
            "composition": assem_result.get("assembly_composition"),
            "interface_count": interface_count,
            "interfaces": interface_dicts,
        }

        assembly_dictionary = {
            "assembly_id": self.assembly_code,
            "pisa_version": "2.0",
            "assembly": assem_dict,
        }

        self._save_to_json({self.pdb_id: assembly_dictionary}, output_json)
        logging.info(
            f"Finished creating assembly-interfaces dictionary in {time() - start} "
            "seconds"
        )
        return interface_count

    def assembly_xml_to_json(self, assembly_xml, output_json, interfaces):
        """Extracts assembly information from PISA xml file and saves it to json file

        Args:
            assembly_xml (str): path to PISA xml file
            output_json (str): path to output json file
        Returns:
            dict: assembly dictionary

        """

        start = time()

        logging.debug("creating simplified assembly dictionary")

        asroot = ET.parse(assembly_xml).getroot()
        assemblies = asroot.iter("asu_complex")
        assem_result = get_assembly_dict(assemblies)
        interface_count = interfaces

        assem_dict = {
            "id": assem_result.get("assembly_id"),
            "size": assem_result.get("assembly_size"),
            "interface_count": interface_count,
            "score": assem_result.get("assembly_score"),
            "macromolecular_size": assem_result.get("assembly_mmsize"),
            "dissociation_energy": assem_result.get("assembly_diss_energy"),
            "accessible_surface_area": assem_result.get("assembly_asa"),
            "buried_surface_area": assem_result.get("assembly_bsa"),
            "entropy": assem_result.get("assembly_entropy"),
            "dissociation_area": assem_result.get("assembly_diss_area"),
            "solvation_energy_gain": assem_result.get("assembly_int_energy"),
            "number_of_uc": assem_result.get("assembly_n_uc"),
            "number_of_dissociated_elements": assem_result.get("assembly_n_diss"),
            "symmetry_number": assem_result.get("assembly_sym_num"),
            "formula": assem_result.get("assembly_formula"),
            "composition": assem_result.get("assembly_composition"),
            "R350": assem_result.get("assembly_R350"),
        }

        assembly_dictionary = {
            "assembly_id": self.assembly_code,
            "pisa_version": "2.0",
            "assembly": assem_dict,
        }

        result = {self.pdb_id: assembly_dictionary}
        self._save_to_json(result, output_json)
        logging.debug(
            f"Finished creating assembly simplified dictionary in {time() - start} "
            "seconds"
        )
        return result

    def _save_to_json(self, data, output_file):
        """Dump the data into a JSON file.

        Args:
            data (dict): The data to be dumped.
            output_file (str): The path to the output file.

        Raises:
            ValueError: If `data` is empty.
        """

        if data:
            with open(output_file, "w") as f:
                json.dump(data, f)
                logging.info(f"Data saved into {output_file} successful")

        else:
            raise ValueError("No data to save")
