#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# IMPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

from __future__ import annotations;

from src.thirdparty.maths import *;
from src.thirdparty.misc import *;
from src.thirdparty.types import *;

from models.generated.commands import *;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# EXPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

__all__ = [
    'Landscape',
];

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# METHOD fitness function -> Array
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class Landscape():
    _fct: np.ndarray;
    _labels: list[str];
    _metric: EnumLandscapeMetric;
    _radius: float;
    _one_based: bool;

    def __init__(
        self,
        values: DataTypeLandscapeValues,
        labels: List[str],
        metric: EnumLandscapeMetric = EnumLandscapeMetric.maximum,
        one_based: bool = False,
    ):
        self._fct = convert_to_nparray(values);
        assert len(labels) == self.dim, 'A label is required for each axis/dimension!';
        self._labels = labels;
        self._metric = metric;
        self._one_based = one_based;
        return;

    @property
    def shape(self) -> tuple:
        return self._fct.shape;

    @property
    def dim(self) -> int:
        return len(self._fct.shape);

    @property
    def coords_middle(self) -> tuple:
        return tuple(math.floor(s/2) for s in self.shape);

    def fitness(self, *x: int) -> float:
        return self._fct[x];

    def label(self, *x: int) -> str:
        if self._one_based:
            x = tuple(xx + 1 for xx in x);
        expr = ','.join([ f'{name}{xx}' for name, xx in zip(self._labels, x)]);
        if self.dim > 1:
            expr = f'({expr})';
        return expr;

    def neighbourhood(self, *x: int, r: float, strict: bool = False) -> List[tuple]:
        r = int(r);
        sides = [
            [ xx - j for j in range(1, r+1) if xx - j in range(s) ]
            + ([ xx ] if xx in range(s) else [])
            + [ xx + j for j in range(1, r+1) if xx + j in range(s) ]
            for xx, s in zip(x, self.shape)
        ];
        match self._metric:
            case EnumLandscapeMetric.maximum:
                umg = list(itertools_product(*sides));
            case EnumLandscapeMetric.manhattan:
                umg = [
                    (*x[:i], xx, *x[(i+1):])
                    for i, side in enumerate(sides)
                    for xx in side
                ];
            case _:
                umg = [ x ];
        if strict:
            umg = [ p for p in umg if p != x ];
        return umg;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# AUXILIARY METHODS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def convert_to_array(values: DataTypeLandscapeValues) -> list:
    return [
        x if isinstance(x, float) else convert_to_array(x)
        for x in values.__root__
    ];

def convert_to_nparray(values: DataTypeLandscapeValues) -> np.ndarray:
    try:
        list_of_lists = convert_to_array(values);
        return np.asarray(list_of_lists, dtype=float);
    except:
        raise ValueError('Could not convert to a d-dimensional array! Ensure that the dimensions are consistent.');