from typing import Any

from mpi4py import MPI
from petsc4py import PETSc

import dolfinx
import numpy as np
import tqdm.autonotebook
import ufl
from dolfinx import fem
from dolfinx.nls.petsc import NewtonSolver
from packaging.version import Version

import festim as F
from festim.mesh.mesh import Mesh as _Mesh
from festim.source import SourceBase as _SourceBase
from festim.subdomain.volume_subdomain import (
    VolumeSubdomain as _VolumeSubdomain,
)


class ProblemBase:
    """
    Base class for :py:class:`HeatTransferProblem <festim.heat_transfer_problem.HeatTransferProblem>` and
    :py:class:`HydrogenTransportProblem <festim.hydrogen_transport_problem.HydrogenTransportProblem>`.

    Attributes:
        show_progress_bar: If `True` a progress bar is displayed during the simulation
        progress_bar: the progress bar
    """

    mesh: _Mesh
    sources: list[_SourceBase]
    exports: list[Any]
    subdomains: list[_VolumeSubdomain]
    show_progress_bar: bool
    progress_bar: None | tqdm.autonotebook.tqdm

    def __init__(
        self,
        mesh: _Mesh = None,
        sources=None,
        exports=None,
        subdomains=None,
        boundary_conditions=None,
        settings=None,
        petsc_options=None,
    ) -> None:
        self.mesh = mesh
        # for arguments to initialise as empty list
        # if arg not None, assign arg, else assign empty list
        self.subdomains = subdomains or []
        self.boundary_conditions = boundary_conditions or []
        self.sources = sources or []
        self.exports = exports or []
        self.settings = settings

        self.dx = None
        self.ds = None
        self.function_space = None
        self.facet_meshtags = None
        self.volume_meshtags = None
        self.formulation = None
        self.bc_forms = []
        self.show_progress_bar = True
        self.petsc_options = petsc_options

    @property
    def volume_subdomains(self):
        return [s for s in self.subdomains if isinstance(s, F.VolumeSubdomain)]

    @property
    def surface_subdomains(self):
        return [s for s in self.subdomains if isinstance(s, F.SurfaceSubdomain)]

    @property
    def dt(self):
        return self._dt

    def define_meshtags_and_measures(self):
        """Defines the facet and volume meshtags of the model which are used
        to define the measures fo the model, dx and ds"""

        if isinstance(self.mesh, F.MeshFromXDMF):
            # TODO: fix naming inconsistency between facet and surface meshtags
            self.facet_meshtags = self.mesh.define_surface_meshtags()
            self.volume_meshtags = self.mesh.define_volume_meshtags()

        elif (
            isinstance(self.mesh, F.Mesh)
            and self.facet_meshtags is None
            and self.volume_meshtags is None
        ):
            self.facet_meshtags, self.volume_meshtags = self.mesh.define_meshtags(
                surface_subdomains=self.surface_subdomains,
                volume_subdomains=self.volume_subdomains,
                # if self has attribute interfaces pass it
                interfaces=getattr(self, "interfaces", None),
            )

        # check volume ids are unique
        vol_ids = [vol.id for vol in self.volume_subdomains]
        if len(vol_ids) != len(np.unique(vol_ids)):
            raise ValueError("Volume ids are not unique")

        # define measures
        self.ds = ufl.Measure(
            "ds", domain=self.mesh.mesh, subdomain_data=self.facet_meshtags
        )
        self.dx = ufl.Measure(
            "dx", domain=self.mesh.mesh, subdomain_data=self.volume_meshtags
        )

    def define_boundary_conditions(self):
        """Defines the dirichlet boundary conditions of the model"""
        for bc in self.boundary_conditions:
            if isinstance(bc, F.DirichletBCBase):
                form = self.create_dirichletbc_form(bc)
                self.bc_forms.append(form)

    def create_solver(self):
        """Creates the solver of the model"""

        if Version(dolfinx.__version__) == Version("0.9.0"):
            problem = fem.petsc.NonlinearProblem(
                self.formulation,
                self.u,
                bcs=self.bc_forms,
            )
            self.solver = NewtonSolver(MPI.COMM_WORLD, problem)

            self.solver.atol = (
                self.settings.atol
                if not callable(self.settings.rtol)
                else self.settings.rtol(float(self.t))
            )
            self.solver.rtol = (
                self.settings.rtol
                if not callable(self.settings.rtol)
                else self.settings.rtol(float(self.t))
            )
            self.solver.max_it = self.settings.max_iterations

            ksp = self.solver.krylov_solver

            if self.petsc_options is None:
                ksp.setType("preonly")
                ksp.getPC().setType("lu")
                ksp.getPC().setFactorSolverType("mumps")
                ksp.setErrorIfNotConverged(True)
            else:
                # Set PETSc options
                opts = PETSc.Options()
                option_prefix = ksp.getOptionsPrefix()
                for k, v in self.petsc_options.items():
                    opts[f"{option_prefix}{k}"] = v
                ksp.setFromOptions()

        elif Version(dolfinx.__version__) > Version("0.9.0"):
            from dolfinx.fem.petsc import NonlinearProblem

            if self.petsc_options is None:
                # taken from https://github.com/FEniCS/dolfinx/blob/5fcb988c5b0f46b8f9183bc844d8f533a2130d6a/python/demo/demo_cahn-hilliard.py#L279C1-L286C28
                use_superlu = (
                    PETSc.IntType == np.int64
                )  # or PETSc.ScalarType == np.complex64
                sys = PETSc.Sys()  # type: ignore
                if sys.hasExternalPackage("mumps") and not use_superlu:
                    linear_solver = "mumps"
                elif sys.hasExternalPackage("superlu_dist"):
                    linear_solver = "superlu_dist"
                else:
                    linear_solver = "petsc"

                petsc_options = {
                    "snes_type": "newtonls",
                    "snes_linesearch_type": "none",
                    "snes_stol": np.sqrt(np.finfo(dolfinx.default_real_type).eps)
                    * 1e-2,
                    # TODO : make atol and rtol callable
                    "snes_atol": self.settings.atol,
                    "snes_rtol": self.settings.rtol,
                    "snes_max_it": self.settings.max_iterations,
                    "ksp_type": "preonly",
                    "pc_type": "lu",
                    "pc_factor_mat_solver_type": linear_solver,
                }
            else:
                petsc_options = self.petsc_options

            self.solver = NonlinearProblem(
                self.formulation,
                self.u,
                bcs=self.bc_forms,
                petsc_options=petsc_options,
                petsc_options_prefix="festim_solver",
            )
            # Delete PETSc options post setting them, ref:
            # https://gitlab.com/petsc/petsc/-/issues/1201
            snes = self.solver.solver
            prefix = snes.getOptionsPrefix()
            opts = PETSc.Options()
            for k in petsc_options.keys():
                del opts[f"{prefix}{k}"]

    def run(self):
        """Runs the model"""

        if self.settings.transient:
            # Solve transient
            if self.show_progress_bar:
                self.progress_bar = tqdm.autonotebook.tqdm(
                    desc=f"Solving {self.__class__.__name__}",
                    total=self.settings.final_time,
                    unit_scale=True,
                )
            while self.t.value < self.settings.final_time:
                self.iterate()
            if self.show_progress_bar:
                self.progress_bar.refresh()  # refresh progress bar to show 100%
                self.progress_bar.close()
        else:
            # Solve steady-state
            if Version(dolfinx.__version__) == Version("0.9.0"):
                self.solver.solve(self.u)
            elif Version(dolfinx.__version__) > Version("0.9.0"):
                self.solver.solve()
            self.post_processing()

    def iterate(self):
        """Iterates the model for a given time step"""
        if self.show_progress_bar:
            self.progress_bar.update(
                min(self.dt.value, abs(self.settings.final_time - self.t.value))
            )

        # update rtol if it's callable
        if callable(self.settings.rtol):
            self.solver.rtol = self.settings.rtol(self.t.value)
        # update rtol if it's callable
        if callable(self.settings.atol):
            self.solver.atol = self.settings.atol(self.t.value)

        self.t.value += self.dt.value

        self.update_time_dependent_values()

        # solve main problem
        if Version(dolfinx.__version__) == Version("0.9.0"):
            nb_its, converged = self.solver.solve(self.u)
        elif Version(dolfinx.__version__) > Version("0.9.0"):
            _ = self.solver.solve()
            converged_reason = self.solver.solver.getConvergedReason()
            assert converged_reason > 0
            nb_its = self.solver.solver.getIterationNumber()

        # post processing
        self.post_processing()

        # update previous solution
        self.u_n.x.array[:] = self.u.x.array[:]

        # adapt stepsize
        if self.settings.stepsize.adaptive:
            new_stepsize = self.settings.stepsize.modify_value(
                value=self.dt.value, nb_iterations=nb_its, t=self.t.value
            )
            self.dt.value = new_stepsize

    def update_time_dependent_values(self):
        t = float(self.t)
        for bc in self.boundary_conditions:
            if bc.time_dependent:
                bc.update(t=t)

        for source in self.sources:
            if source.value.explicit_time_dependent:
                source.value.update(t=t)
