# genotype.py
import json
import math
import random
import time
from typing import Dict, List, Tuple, Any, Optional


def generate_id() -> float:
    return random.random()


def create_neural_weights(vector_length: int) -> List[float]:
    return [random.uniform(-2.0, 2.0) for _ in range(vector_length)]


def construct(
        morphology_module,
        hidden_layer_densities: List[int],
        file_name: Optional[str] = None,
        *,
        add_bias: bool = False,
) -> Dict[str, Any]:
    rnd_seed = time.time_ns() & 0xFFFFFFFF
    random.seed(rnd_seed)

    S = morphology_module.get_InitSensor(morphology_module)
    A = morphology_module.get_InitActuator(morphology_module)

    sensor = {
        "id": S.get("id", generate_id()),
        "name": S["name"],
        "vector_length": int(S["vector_length"]),
        "cx_id": None,  # wird später gesetzt
        "fanout_ids": [],  # wird später gesetzt
        # optional:
        # "scape": S.get("scape")
    }

    actuator = {
        "id": A.get("id", generate_id()),
        "name": A["name"],
        "vector_length": int(A["vector_length"]),
        "cx_id": None,  # wird später gesetzt
        "fanin_ids": [],  # wird später gesetzt
        # optional:
        # "scape": A.get("scape")
    }

    output_vl = actuator["vector_length"]
    layer_densities = list(hidden_layer_densities) + [output_vl]

    cortex_id = generate_id()

    layers = _create_neuro_layers(
        cx_id=cortex_id,
        sensor=sensor,
        actuator=actuator,
        layer_densities=layer_densities,
        add_bias=add_bias,
    )

    input_layer = layers[0]
    output_layer = layers[-1]

    sensor["cx_id"] = cortex_id
    sensor["fanout_ids"] = [n["id"] for n in input_layer]

    actuator["cx_id"] = cortex_id
    actuator["fanin_ids"] = [n["id"] for n in output_layer]

    neuron_ids = [n["id"] for layer in layers for n in layer]
    cortex = {
        "id": cortex_id,
        "sensor_ids": [sensor["id"]],
        "actuator_ids": [actuator["id"]],
        "neuron_ids": neuron_ids,
    }

    # 7) Genotyp zusammensetzen
    genotype = {
        "cortex": cortex,
        "sensor": sensor,
        "actuator": actuator,
        "neurons": [n for layer in layers for n in layer],
    }

    # 8) Optional speichern
    if file_name:
        save_genotype(file_name, genotype)

    return genotype


def _create_neuro_layers(
        cx_id: float,
        sensor: Dict[str, Any],
        actuator: Dict[str, Any],
        layer_densities: List[int],
        *,
        add_bias: bool,
) -> List[List[Dict[str, Any]]]:
    layers: List[List[Dict[str, Any]]] = []

    input_idps: List[Tuple[float, int]] = [(sensor["id"], sensor["vector_length"])]

    for layer_index, layer_density in enumerate(layer_densities):
        neuron_ids = [generate_id() for _ in range(layer_density)]

        if layer_index < len(layer_densities) - 1:
            next_ids = [generate_id() for _ in range(layer_densities[layer_index + 1])]
            output_ids = next_ids
        else:
            output_ids = [actuator["id"]]

        this_layer: List[Dict[str, Any]] = []
        for _nid in neuron_ids:
            proper_input = _create_neural_input(input_idps, add_bias=add_bias)
            neuron = {
                "id": _nid,
                "layer_index": layer_index,
                "cx_id": cx_id,
                "activation_function": "tanh",
                "input_weights": [{"input_id": i, "weights": w} for (i, w) in proper_input],
                "output_ids": output_ids[:],  # Kopie
            }
            this_layer.append(neuron)

        layers.append(this_layer)
        input_idps = [(n["id"], 1) for n in this_layer]

    return layers


def _is_bias_tuple(t: Tuple[Any, Any]) -> bool:
    key, _ = t
    return isinstance(key, str) and key == "bias"


def _create_neural_input(
        input_idps: List[Tuple[float, int]],
        *,
        add_bias: bool,
) -> List[Tuple[Any, List[float]]]:
    proper: List[Tuple[Any, List[float]]] = []
    for input_id, vl in input_idps:
        proper.append((input_id, create_neural_weights(vl)))

    if add_bias:
        proper.append(("bias", [random.random() - 0.5]))

    return proper


def save_genotype(file_name: str, genotype: Dict[str, Any]) -> None:
    with open(file_name, "w") as f:
        json.dump(genotype, f, indent=2)


def load_from_file(file_name: str) -> Dict[str, Any]:
    with open(file_name, "r") as f:
        return json.load(f)


def print_genotype(file_name: str) -> None:
    g = load_from_file(file_name)
    cx = g["cortex"]
    print("[CORTEX]", cx)
    sids = cx.get("sensor_ids", [])
    nids = cx.get("neuron_ids", [])
    aids = cx.get("actuator_ids", [])

    nid2n = {n["id"]: n for n in g.get("neurons", [])}
    sid2s = {g["sensor"]["id"]: g["sensor"]} if "sensor" in g else {s["id"]: s for s in g.get("sensors", [])}
    aid2a = {g["actuator"]["id"]: g["actuator"]} if "actuator" in g else {a["id"]: a for a in g.get("actuators", [])}

    for sid in sids:
        print("[SENSOR]", sid2s.get(sid))
    for nid in nids:
        print("[NEURON]", nid2n.get(nid))
    for aid in aids:
        print("[ACTUATOR]", aid2a.get(aid))