from __future__ import annotations

import csv
import os
from collections import defaultdict
from inspect import Parameter, Signature, cleandoc, signature
from typing import Literal

from numpydoc.docscrape import Parameter as numpydoc_Parameter
from numpydoc.docscrape_sphinx import SphinxDocString

from dpdata.bond_order_system import BondOrderSystem

# ensure all plugins are loaded!
from dpdata.driver import Driver, Minimizer
from dpdata.format import Format
from dpdata.system import LabeledSystem, MultiSystems, System


def get_formats() -> dict:
 formats = defaultdict(list)
 for kk, ff in Format.get_formats().items():
 formats[ff].append(kk)
 return formats


def get_driver() -> dict:
 drivers = defaultdict(list)
 for kk, ff in Driver.get_drivers().items():
 drivers[ff].append(kk)
 return drivers


def get_minimizer() -> dict:
 minimizers = defaultdict(list)
 for kk, ff in Minimizer.get_minimizers().items():
 minimizers[ff].append(kk)
 return minimizers


def detect_overridden(cls: Format, method: str) -> bool:
 """Check whether a method is override.

 Parameters
 ----------
 cls : Format
 a format
 method : str
 method name

 Returns
 -------
 bool
 whether a method is overridden
 """
 return method in cls.__dict__


def get_cls_link(cls: object) -> str:
 """Returns class link.

 Parameters
 ----------
 cls : object
 the class

 Returns
 -------
 str
 the link of a class
 """
 return ":class:`{} <{}>`".format(
 cls.__name__, ".".join([cls.__module__, cls.__name__])
 )


def check_supported(fmt: Format):
 methods = {}
 for mtd in [
 "from_system",
 "to_system",
 "from_labeled_system",
 "to_labeled_system",
 "from_bond_order_system",
 "to_bond_order_system",
 "from_multi_systems",
 "to_multi_systems",
 ]:
 if detect_overridden(fmt, mtd):
 methods[mtd] = None
 if mtd == "to_system":
 methods["to_labeled_system"] = None
 if fmt.MultiMode != fmt.MultiModes.NotImplemented:
 methods["from_multi_systems"] = None
 methods["to_multi_systems"] = None
 return list(methods.keys())


method_links = {
 "from_system": ":ref:`System() <{}_{}>`",
 "to_system": ":ref:`System.to() <{}_{}>`",
 "from_labeled_system": ":ref:`LabeledSystem() <{}_{}>`",
 "to_labeled_system": ":ref:`LabeledSystem.to() <{}_{}>`",
 "from_bond_order_system": ":ref:`BondOrderSystem() <{}_{}>`",
 "to_bond_order_system": ":ref:`BondOrderSystem.to() <{}_{}>`",
 "from_multi_systems": ":ref:`MultiSystems.load_systems_from_file() <{}_{}>`",
 "to_multi_systems": ":ref:`MultiSystems.to() <{}_{}>`",
}

method_classes = {
 "from_system": "System",
 "to_system": "System",
 "from_labeled_system": "LabeledSystem",
 "to_labeled_system": "LabeledSystem",
 "from_bond_order_system": "BondOrderSystem",
 "to_bond_order_system": "BondOrderSystem",
 "from_multi_systems": "MultiSystems",
 "to_multi_systems": "MultiSystems",
}

method_cls_obj = {
 "from_system": System,
 "to_system": System,
 "from_labeled_system": LabeledSystem,
 "to_labeled_system": LabeledSystem,
 "from_bond_order_system": BondOrderSystem,
 "to_bond_order_system": BondOrderSystem,
 "from_multi_systems": MultiSystems,
 "to_multi_systems": MultiSystems,
}


def generate_sub_format_pages(formats: dict):
 """Generate sub format pages."""
 os.makedirs("formats", exist_ok=True)
 for format, alias in formats.items():
 # format: Format, alias: list[str]
 buff = []
 buff.append(f".. _{format.__name__}:")
 buff.append("")
 for aa in alias:
 buff.append(f"{aa} format")
 buff.append("=" * len(buff[-1]))
 buff.append("")
 buff.append(f"Class: {get_cls_link(format)}")
 buff.append("")

 docstring = format.__doc__
 if docstring is not None:
 docstring = cleandoc(docstring)
 rst = str(SphinxDocString(docstring))
 buff.append(rst)
 buff.append("")

 buff.append("Conversions")
 buff.append("-----------")
 methods = check_supported(format)
 for method in methods:
 buff.append("")
 buff.append(f".. _{format.__name__}_{method}:")
 buff.append("")
 if method.startswith("from_"):
 buff.append(f"Convert from this format to {method_classes[method]}")
 buff.append("`" * len(buff[-1]))
 elif method.startswith("to_"):
 buff.append(f"Convert from {method_classes[method]} to this format")
 buff.append("`" * len(buff[-1]))
 buff.append("")
 method_obj = getattr(format, method)
 if (
 method == "to_labeled_system"
 and method not in format.__dict__
 and "to_system" in format.__dict__
 ):
 method_obj = getattr(format, "to_system")
 docstring = method_obj.__doc__
 if docstring is not None:
 docstring = cleandoc(docstring)
 sig = signature(method_obj)
 parameters = dict(sig.parameters)
 return_annotation = sig.return_annotation
 # del self
 del parameters[list(parameters)[0]]
 # del data
 if method.startswith("to_"):
 del parameters[list(parameters)[0]]
 if "args" in parameters:
 del parameters["args"]
 if "kwargs" in parameters:
 del parameters["kwargs"]
 if method == "to_multi_systems" or method.startswith("from_"):
 sig = Signature(
 list(parameters.values()), return_annotation=method_cls_obj[method]
 )
 else:
 sig = Signature(
 list(parameters.values()), return_annotation=return_annotation
 )
 sig = str(sig)
 if method.startswith("from_"):
 if method != "from_multi_systems":
 for aa in alias:
 parameters["fmt"] = Parameter(
 "fmt",
 Parameter.POSITIONAL_OR_KEYWORD,
 default=None,
 annotation=Literal[aa],
 )
 sig_fmt = Signature(
 list(parameters.values()),
 return_annotation=method_cls_obj[method],
 )
 sig_fmt = str(sig_fmt)
 buff.append(
 f""".. py:function:: dpdata.{method_classes[method]}{sig_fmt}"""
 )
 buff.append(""" :noindex:""")
 for aa in alias:
 buff.append(
 """.. py:function:: dpdata.{}.from_{}{}""".format(
 method_classes[method],
 aa.replace("/", "_").replace(".", ""),
 sig,
 )
 )
 buff.append(""" :noindex:""")
 buff.append("")
 if docstring is None or method not in format.__dict__:
 docstring = f""" Convert this format to :class:`{method_classes[method]}`."""
 doc_obj = SphinxDocString(docstring)
 if len(doc_obj["Parameters"]) > 0:
 doc_obj["Parameters"] = [
 xx
 for xx in doc_obj["Parameters"]
 if xx.name not in ("*args", "**kwargs")
 ]
 else:
 if method == "from_multi_systems":
 doc_obj["Parameters"] = [
 numpydoc_Parameter(
 "directory",
 "str",
 ["directory of systems"],
 )
 ]
 doc_obj["Yields"] = []
 doc_obj["Returns"] = [
 numpydoc_Parameter("", method_classes[method], ["converted system"])
 ]
 rst = " " + str(doc_obj)
 buff.append(rst)
 buff.append("")
 elif method.startswith("to_"):
 for aa in alias:
 parameters = {
 "fmt": Parameter(
 "fmt",
 Parameter.POSITIONAL_OR_KEYWORD,
 annotation=Literal[aa],
 ),
 **parameters,
 }
 if method == "to_multi_systems":
 sig_fmt = Signature(
 list(parameters.values()),
 return_annotation=method_cls_obj[method],
 )
 else:
 sig_fmt = Signature(
 list(parameters.values()),
 return_annotation=return_annotation,
 )
 sig_fmt = str(sig_fmt)
 buff.append(
 f""".. py:function:: dpdata.{method_classes[method]}.to{sig_fmt}"""
 )
 buff.append(""" :noindex:""")
 for aa in alias:
 buff.append(
 """.. py:function:: dpdata.{}.to_{}{}""".format(
 method_classes[method],
 aa.replace("/", "_").replace(".", ""),
 sig,
 )
 )
 buff.append(""" :noindex:""")
 buff.append("")
 if docstring is None or (
 method not in format.__dict__
 and not (
 method == "to_labeled_system"
 and method not in format.__dict__
 and "to_system" in format.__dict__
 )
 ):
 docstring = (
 f"Convert :class:`{method_classes[method]}` to this format."
 )
 doc_obj = SphinxDocString(docstring)
 if len(doc_obj["Parameters"]) > 0:
 doc_obj["Parameters"] = [
 xx
 for xx in doc_obj["Parameters"][1:]
 if xx.name not in ("*args", "**kwargs")
 ]
 else:
 if method == "to_multi_systems":
 doc_obj["Parameters"] = [
 numpydoc_Parameter(
 "directory",
 "str",
 ["directory to save systems"],
 )
 ]
 if method == "to_multi_systems":
 doc_obj["Yields"] = []
 doc_obj["Returns"] = [
 numpydoc_Parameter("", method_classes[method], ["this system"])
 ]
 rst = " " + str(doc_obj)
 buff.append(rst)
 buff.append("")
 buff.append("")
 buff.append("")

 with open(f"formats/{format.__name__}.rst", "w") as rstfile:
 rstfile.write("\n".join(buff))


if __name__ == "__main__":
 formats = get_formats()
 with open("formats.csv", "w", newline="") as csvfile:
 fieldnames = [
 "Format",
 "Alias",
 "Supported Conversions",
 ]
 writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

 writer.writeheader()
 for kk, vv in formats.items():
 writer.writerow(
 {
 "Format": f":ref:`{kk.__name__}`",
 "Alias": "\n".join(f"``{vvv}``" for vvv in vv),
 "Supported Conversions": "\n".join(
 method_links[mtd].format(kk.__name__, mtd)
 for mtd in check_supported(kk)
 ),
 }
 )

 drivers = get_driver()
 with open("drivers.csv", "w", newline="") as csvfile:
 fieldnames = [
 "Class",
 "Alias",
 ]
 writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

 writer.writeheader()
 for kk, vv in drivers.items():
 writer.writerow(
 {
 "Class": get_cls_link(kk),
 "Alias": "\n".join(f"``{vvv}``" for vvv in vv),
 }
 )

 minimizers = get_minimizer()
 with open("minimizers.csv", "w", newline="") as csvfile:
 fieldnames = [
 "Class",
 "Alias",
 ]
 writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

 writer.writeheader()
 for kk, vv in minimizers.items():
 writer.writerow(
 {
 "Class": get_cls_link(kk),
 "Alias": "\n".join(f"``{vvv}``" for vvv in vv),
 }
 )
 generate_sub_format_pages(formats)