# Copyright (C) 2023 - 2026 ANSYS, Inc. and/or its affiliates.
# SPDX-License-Identifier: MIT
#
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import logging
from typing import Dict, List, Literal, Optional

from tqdm import tqdm
from wakepy import keep

from ansys.simai.core.data.base import ComputableDataModel, Directory
from ansys.simai.core.data.geometries import Geometry
from ansys.simai.core.data.types import (
    Identifiable,
    get_id_from_identifiable,
    get_object_from_identifiable,
)
from ansys.simai.core.errors import InvalidArguments

logger = logging.getLogger(__name__)


class _OptimizationTrialRun(ComputableDataModel):
    """Provides the local representation of an optimization trial run object.

    The optimization trial run is an iteration of the optimization process.
    Each trial run tests a geometry and returns new parameters for the next geometry to try.
    """


# Undocumented for now, users don't really need to interact with it
class _OptimizationTrialRunDirectory(Directory[_OptimizationTrialRun]):
    _data_model = _OptimizationTrialRun

    def get(self, trial_run_id: str):
        """Get a specific trial run from the server."""
        return self._model_from(self._client._api.get_optimization_trial_run(trial_run_id))

    def _run_iteration(
        self, optimization: Identifiable["Optimization"], parameters: Dict
    ) -> _OptimizationTrialRun:
        optimization_id = get_id_from_identifiable(optimization)
        return self._model_from(
            self._client._api.run_optimization_trial(optimization_id, parameters)
        )


class Optimization(ComputableDataModel):
    """Provides the local representation of an optimization definition object."""

    def _run_iteration(self, parameters: Dict) -> "_OptimizationTrialRun":
        return self._client._optimization_trial_run_directory._run_iteration(self.id, parameters)


class OptimizationResult:
    """Result for a non-parametric optimization."""

    def __init__(
        self,
        optimization: "Optimization",
        objectives: list[Dict],
        geometries: list[Dict],
    ):
        self._optimization = optimization
        self._objectives: list[Dict] = objectives
        self._geometries: list[Geometry] = [
            self._optimization._client.geometries._model_from(geom) for geom in geometries
        ]

    @property
    def optimization(self) -> Optimization:
        """Optimization object of the optimization run."""
        return self._optimization

    def list_objectives(self) -> List[Dict]:
        """List of objectives generated by the optimization."""
        return self._objectives

    def list_geometries(self) -> List[Geometry]:
        """List of geometries generated by the optimization."""
        return self._geometries


class OptimizationDirectory(Directory[Optimization]):
    """Provides a collection of methods related to optimizations.

    This class is accessed through ``client.optimizations``.

    Example:
        .. code-block:: python

            import ansys.simai.core as asc

            simai_client = asc.from_config()
            simai_client.optimizations.run_non_parametric(...)
    """

    _data_model = Optimization

    def get(self, optimization_id: str) -> Optimization:
        """Get a specific optimization object from the server.

        Args:
            optimization_id: ID of the optimization.

        Returns:
            :py:class:`Optimization`.
        """
        return self._model_from(self._client._api.get_optimization(optimization_id))

    def run_non_parametric(
        self,
        geometry: Identifiable[Geometry],
        bounding_boxes: List[List[float]],
        n_iters: int,
        symmetries: Optional[List[Literal["x", "y", "z", "X", "Y", "Z"]]] = None,
        axial_symmetry: Optional[Literal["x", "y", "z"]] = None,
        scalars: Optional[Dict[str, float]] = None,
        minimize: Optional[List[str]] = None,
        maximize: Optional[List[str]] = None,
        max_displacement: Optional[List[float]] = None,
        show_progress: bool = False,
        boundary_conditions: Optional[Dict[str, float]] = None,
    ) -> OptimizationResult:
        """Run an optimization loop to generate geometries, server-side, using automorphing.
        Automorphing is a non-parametric deformation of a surface geometry.

        Args:
            geometry: Required. The object (Geometry) or the ID (str) of the baseline geometry on which to perform
                the non-parametric optimization and that has already been used to build an AI model.
                The optimization will run in the same workspace as the model.
            bounding_boxes: Required. The list of the bounds of the different boxes that will define the absolute locations
                of the geometry to optimize. It is a list of lists, and each sub-list must have exactly six items with the following
                expected order: ``[x_min, x_max, y_min, y_max, z_min, z_max]``.

                Example format::

                    [
                    [box1_xmin, box1_xmax, box1_ymin, box1_ymax, box1_zmin, box1_zmax],
                    [box2_xmin, box2_xmax, box2_ymin, box2_ymax, box2_zmin, box2_zmax],
                    ...
                    ]

            n_iters: Required. The number of optimization iterations. This number must be a strictly positive integer.
                It will define the number of automorphed geometries uploaded to the SimAI workspace.
            symmetries: Optional. The list of symmetry axes, axes being x, y, and z, defining a plane around which the geometry is mirrored.

                - The planar symmetry is applied to all the ``bounding_boxes`` defined.
                - ``symmetries`` and ``axial_symmetry`` are mutually exclusive parameters.

            axial_symmetry: Optional. The axis, defined by a scalar, along which axial symmetry is applied.
                For the deformation to be axially symmetrical:

                - The center of the bounding box must coincide with the targeted central axis.
                - The mesh should be axially symmetrical in the bounding box.
                - The ``axial_symmetry`` is applied to all the ``bounding_boxes`` defined.
                - ``symmetries`` and ``axial_symmetry`` are mutually exclusive parameters.

            scalars: Optional. The values of the scalars to perform the optimization at.
                The values must correspond to existing scalars already defined in your SimAI workspace.
            minimize: Required if no ``maximize`` parameter is defined. A list of one global coefficient to minimize.
                This global coefficient must correspond to one of the existing coefficients defined in your model configuration.
            maximize: Required if no ``minimize`` parameter is defined. A list of one global coefficient to maximize.
                This global coefficient must correspond to one of the existing coefficients defined in your model configuration.

                .. note::

                    - ``minimize`` and ``maximize`` are mutually exclusive objectives; define only one.
                    - The defined objective must be computed from the surface fields because mesh nodes must be involved.

            max_displacement: Optional. User-defined constraint on the maximum allowable deformation of the initial mesh in non-parametric optimization.
                It is specified as a list (``max_displacement``) matching the number of bounding boxes (``bounding_boxes``).

                For example, for two bounding boxes:

                - ``bounding_boxes = [[0,1,0,2,0,4],[10,2,10,4,10,5]]``
                - ``max_displacement = [0.002, 0.001]``

                Each value limits the displacement within the corresponding bounding box, using the same metric as the bounding box coordinates.
            show_progress: Optional. Whether to print progress bar on stdout.
                It is updated each time a new iteration is completed.
            boundary_conditions: Optional. **(Deprecated)** The values of the boundary conditions to perform the optimization at.
                The values must correspond to existing boundary conditions already defined in your SimAI workspace.

        Example:
          .. code-block:: python

            import ansys.simai.core as asc

            simai_client = asc.from_config(workspace="optimization-workspace")
            geometry = simai_client.geometries.list()[0]

            simai_client.optimizations.run_non_parametric(
                geometry,
                bounding_boxes=[[0, 1, 0, 1, 0, 1]],
                scalars={"VelocityX": 10.5},
                symmetries=["y"],
                n_iters=10,
                minimize=["TotalForceX"],
                show_progress=True,
            )

        Returns:
            An object containing the results of the optimization.

        """
        if scalars is None and boundary_conditions is not None:
            logger.warning(
                "The 'boundary_conditions' parameter is deprecated and will be removed in a future release. Please use the 'scalars' parameter instead."
            )
            scalars = boundary_conditions
        _validate_n_iters(n_iters)
        _validate_global_coefficients_for_non_parametric(minimize, maximize)
        _validate_bounding_boxes(bounding_boxes)
        _validate_max_displacement(max_displacement, bounding_boxes)
        _validate_axial_symmetry(axial_symmetry, symmetries)
        geometry = get_object_from_identifiable(geometry, self._client._geometry_directory)
        objective = _build_objective(minimize, maximize)
        optimization_parameters = {
            "boundary_conditions": scalars or {},
            "n_iters": n_iters,
            "objective": objective,
            "type": "non_parametric",
            "geometry_generation": {
                "geometry": geometry.id,
                "box_bounds_list": bounding_boxes,
                "symmetries": symmetries or [],
                "max_displacement": max_displacement,
                "axial_symmetry": axial_symmetry,
            },
        }
        with tqdm(total=n_iters, disable=not show_progress) as progress_bar:
            progress_bar.set_description("Creating optimization definition.")
            optimization = self._model_from(
                self._client._api.define_optimization(
                    geometry._fields["workspace_id"], optimization_parameters
                )
            )
            optimization.wait()
            trial_result = {"objectives": [], "geometries": []}
            logger.debug("Optimization defined. Starting optimization loop.")
            with keep.running(on_fail="warn"):
                for _ in range(n_iters):
                    logger.debug("Running iteration")
                    progress_bar.set_description("Running iteration")
                    trial_run = optimization._run_iteration({})
                    trial_run.wait()
                    progress_bar.set_postfix(
                        {
                            "objective": trial_run.fields["outcome_values"],
                        }
                    )
                    if trial_run.fields.get("is_feasible", True):
                        trial_result["objectives"].append(trial_run.fields["outcome_values"])
                        trial_result["geometries"].append(trial_run.fields["geometry"])
                    else:
                        logger.debug("Trial run results did not match constraints. Skipping.")
                    logger.debug("Trial complete.")
                    progress_bar.update(1)
                logger.debug("Optimization complete.")
                progress_bar.set_description("Optimization complete.")

            optimization_result = OptimizationResult(
                optimization, trial_result["objectives"], trial_result["geometries"]
            )
            return optimization_result


def _validate_bounding_boxes(bounding_boxes: list[list[float]]) -> None:
    if not isinstance(bounding_boxes, list) or not bounding_boxes:
        raise InvalidArguments("bounding_boxes must be a non-empty list.")

    for i, box in enumerate(bounding_boxes):
        if not isinstance(box, list) or len(box) != 6:
            raise InvalidArguments(
                f"Bounding box at index {i} should be a list of 6 values "
                f"[xmin, xmax, ymin, ymax, zmin, zmax]."
            )

        if not all(isinstance(value, (int, float)) for value in box):
            raise InvalidArguments(f"Bounding box at index {i} contains non-numeric values.")

        dimensions = ("x", "y", "z")
        for j in range(3):
            min_index = j * 2
            max_index = min_index + 1
            if not (box[min_index] < box[max_index]):
                raise InvalidArguments(
                    f"Bounding box at index {i}: {dimensions[j]}min ({box[min_index]}) "
                    f"must be less than {dimensions[j]}max ({box[max_index]})."
                )


def _validate_global_coefficients_for_non_parametric(minimize: List[str], maximize: List[str]):
    if minimize and maximize:
        raise InvalidArguments(
            "Only one of minimize or maximize can be provided for non parametric optimization."
        )

    active_list = minimize if minimize else maximize
    if not isinstance(active_list, list) or (
        isinstance(active_list, list) and len(active_list) != 1
    ):
        raise InvalidArguments(
            "minimize or maximize must be a list of one string for non parametric optimization."
        )


def _validate_outcome_constraints(outcome_constraints: list) -> None:
    if not isinstance(outcome_constraints, list):
        raise InvalidArguments("outcome_constraints must be a list")

    for i, constraint in enumerate(outcome_constraints):
        if not isinstance(constraint, str):
            raise InvalidArguments(f"Constraint at index {i} must be a string")

        # Check if constraint contains either >= or <=
        if ">=" in constraint:
            operator = ">="
        elif "<=" in constraint:
            operator = "<="
        else:
            raise InvalidArguments(
                f"Constraint at index {i} ({constraint}) must contain either >= or <= operator"
            )

        parts = constraint.split(operator)
        if len(parts) != 2:
            raise InvalidArguments(
                f"Constraint at index {i} ({constraint}) must be of form 'metric_name {operator} value'"
            )

        metric_name = parts[0].strip()
        value = parts[1].strip()

        if not metric_name:
            raise InvalidArguments(
                f"Constraint at index {i} ({constraint}) must have a metric name"
            )

        try:
            float(value)
        except ValueError as error:
            raise InvalidArguments(
                f"Constraint at index {i} ({constraint}): '{value}' must be a numeric value"
            ) from error


def _validate_n_iters(n_iters) -> None:
    if not isinstance(n_iters, int):
        raise InvalidArguments("n_iters must be an integer")

    if n_iters <= 0:
        raise InvalidArguments("n_iters must be strictly positive")


def _validate_max_displacement(
    max_displacement: Optional[List[float]], bounding_boxes: List[List[float]]
) -> None:
    if max_displacement is None:
        return

    if not isinstance(max_displacement, list):
        raise InvalidArguments("max_displacement must be a list")

    if not all(isinstance(value, (int, float)) for value in max_displacement):
        raise InvalidArguments("max_displacement contains non-numeric values")

    if len(max_displacement) != len(bounding_boxes):
        raise InvalidArguments(
            "Max displacement list and bounding boxes list must have the same number of items"
        )


def _validate_axial_symmetry(
    axial_symmetry: Optional[Literal["x", "y", "z"]],
    symmetries: Optional[List[Literal["x", "y", "z", "X", "Y", "Z"]]],
) -> None:
    if axial_symmetry is not None and axial_symmetry not in ("x", "y", "z"):
        raise InvalidArguments("axial_symmetry must be one of 'x', 'y', 'z' or None")

    if axial_symmetry is not None and symmetries is not None:
        raise InvalidArguments("symmetries and axial_symmetry are mutually exclusive parameters")


def _build_objective(minimize: list[str], maximize: list[str]) -> dict:
    if not minimize and not maximize:
        raise InvalidArguments("No global coefficient to optimize.")
    objective = {}
    if minimize:
        for global_coefficient in minimize:
            objective[global_coefficient] = {"minimize": True}
    if maximize:
        for global_coefficient in maximize:
            objective[global_coefficient] = {"minimize": False}
    return objective
