from pathlib import Path
import typer
from typing import Annotated, Optional

from ...helper.flux import machines, flux
from ...helper.experiment import ExperimentConfig, LogViewType


_base_side_length = 0.1
_efficiency_base_config = f"""# Single material, simple test to determine optimal point/GPU
method: mpm
mpm:
  point_location_type: uniform

material: cube
cube:
  model: elasticity-mpm-neo-hookean-damage-current
  E: 50e2
  nu: 0.4
  rho: 15
  fracture_toughness: 1e3 # increase to reduce damage accumulation
  characteristic_length: 1 # 4*mesh_cl
  residual_stiffness: 1e-2
  damage_scaling: 1
  use_AT1: false


dm_plex:
  # Mesh specification determined by script
  dim: 3
  simplex: 0
  # Use spatial hashing for point location
  hash_location:

bc:
  # Clamp all sides
  clamp: 1,2,3,4,5,6

forcing: body
forcing_body_vector: -9.81,0,0

# TS
ts:
  dt: 0.001
  max_time: 0.01
  monitor:

# SNES
snes:
  type: newtonls
  linesearch:
    type: cp
    max_it: 10

# KSP
ksp:
  type: bcgsl
  norm_type: preconditioned

# Preconditioner
pc:
  type: jacobi
"""


class EfficiencyForcingExperiment(ExperimentConfig):
    """Efficiency experiment configurations for optimal point/GPU determination, no mesh remapping."""
    num_cells_1d: int

    def __init__(self, num_cells_1d: int, order: int = 1, points_per_cell: int = 27):
        self.num_cells_1d = num_cells_1d
        self.num_points_per_cell = points_per_cell
        self.order = order
        super().__init__(Path(__file__).stem.replace('_', '-'), self.__doc__, _efficiency_base_config)

    @property
    def mesh_options(self) -> str:
        # Scale the mesh size based on the number of faces
        side_length = _base_side_length * self.num_cells_1d
        mesh_options = '\n'.join([
            "",
            "# Experiment options for efficiency experiment",
            f"mpm_num_points_per_cell: {self.num_points_per_cell}",
            f"order: {self.order}",
            "",
            "# Mesh options generated by script:",
            "dm_plex:",
            f"  box_faces: {self.num_cells_1d},{self.num_cells_1d},{self.num_cells_1d}",
            f"  box_upper: {side_length},{side_length},{side_length}",
            "",
        ])
        return mesh_options

    @property
    def name(self) -> str:
        return f"{self._name}_{self.num_cells_1d}"

    def __str__(self) -> str:
        output = '\n'.join([
            f'[h1]Ratel iMPM Efficiency Experiment[/]',
            f'{self.description}',
            f"\n[h2]Mesh Options[/]",
            f"  • Solution basis order: {self.order}",
            f"  • Number of cells: {self.num_cells_1d}^3={self.num_cells_1d**3}",
            f"  • Points per cell: {self.num_points_per_cell}",
            f"  • Points per GPU: {self.num_points_per_cell * self.num_cells_1d**3}",
        ])
        if self.user_options:
            output += "\n[h2]User Options[/]\n"
            output += "\n".join([f"  • {key}: {value}" for key, value in self.user_options.items()])
        return output


__doc__ = EfficiencyForcingExperiment.__doc__

app = typer.Typer(name="efficiency", help=__doc__)


@app.command()
def write_config(num_cells_1d: Annotated[int, typer.Argument(min=3)],
                 output_dir: Path, log_view: Optional[LogViewType] = None):
    """Generate the efficiency experiment configuration."""
    experiment = EfficiencyForcingExperiment(num_cells_1d)
    experiment.logview = log_view
    experiment.write_config(output_dir)


@app.command(
    context_settings={"allow_extra_args": True, "ignore_unknown_options": True}
)
def flux_run(
    ctx: typer.Context,
    num_cells_1d: Annotated[int, typer.Argument(min=3)],
    order: Annotated[int, typer.Argument(min=1)] = 1,
    points_per_cell: Annotated[int, typer.Argument(min=8)] = 27,
    num_processes: Annotated[int, typer.Option("-n", min=1)] = 1,
    max_time: Annotated[str, typer.Option('-t')] = '1h',
    log_view: Optional[LogViewType] = None,
    machine: Optional[machines.Machine] = None,
    dry_run: bool = False
):
    """Run the efficiency experiment using flux."""
    experiment = EfficiencyForcingExperiment(num_cells_1d, order, points_per_cell)
    experiment.user_options = ctx.args
    experiment.logview = log_view
    script_file, _ = flux.generate(
        experiment,
        machine=machine,
        num_processes=num_processes,
        max_time=max_time,
    )
    if not dry_run:
        flux.run(script_file)
