from __future__ import annotations

from typing import Dict, List, Sequence


MOVEMENT_STRONG_DOWN_THRESHOLD = -0.40
MOVEMENT_DOWN_THRESHOLD = -0.30
MOVEMENT_PUMP_50_THRESHOLD = 0.50
MOVEMENT_PUMP_100_THRESHOLD = 1.00
MOVEMENT_PUMP_300_THRESHOLD = 3.00

MOVEMENT_CLASS_NAMES = [
    "strong_down",
    "down",
    "flat",
    "up",
    "strong_up",
    "extreme_up",
]
MOVEMENT_CLASS_TO_ID = {name: idx for idx, name in enumerate(MOVEMENT_CLASS_NAMES)}
MOVEMENT_ID_TO_CLASS = {idx: name for name, idx in MOVEMENT_CLASS_TO_ID.items()}

DEFAULT_MOVEMENT_LABEL_CONFIG = {
    "strong_down_threshold": MOVEMENT_STRONG_DOWN_THRESHOLD,
    "down_threshold": MOVEMENT_DOWN_THRESHOLD,
    "pump_50_threshold": MOVEMENT_PUMP_50_THRESHOLD,
    "pump_100_threshold": MOVEMENT_PUMP_100_THRESHOLD,
    "pump_300_threshold": MOVEMENT_PUMP_300_THRESHOLD,
}


def classify_movement_return(
    return_value: float,
    movement_label_config: Dict[str, float] | None = None,
) -> int:
    cfg = dict(DEFAULT_MOVEMENT_LABEL_CONFIG)
    if movement_label_config:
        cfg.update({k: float(v) for k, v in movement_label_config.items() if k in cfg})

    strong_down_threshold = min(cfg["strong_down_threshold"], cfg["down_threshold"])
    down_threshold = cfg["down_threshold"]
    pump_50_threshold = cfg["pump_50_threshold"]
    pump_100_threshold = cfg["pump_100_threshold"]
    pump_300_threshold = cfg["pump_300_threshold"]

    if return_value <= strong_down_threshold:
        return MOVEMENT_CLASS_TO_ID["strong_down"]
    if return_value < down_threshold:
        return MOVEMENT_CLASS_TO_ID["down"]
    if return_value < pump_50_threshold:
        return MOVEMENT_CLASS_TO_ID["flat"]
    if return_value < pump_100_threshold:
        return MOVEMENT_CLASS_TO_ID["up"]
    if return_value < pump_300_threshold:
        return MOVEMENT_CLASS_TO_ID["strong_up"]
    return MOVEMENT_CLASS_TO_ID["extreme_up"]


def derive_movement_targets(
    horizon_returns: Sequence[float],
    horizon_mask: Sequence[float],
    movement_label_config: Dict[str, float] | None = None,
) -> Dict[str, List[int]]:
    class_targets: List[int] = []
    class_mask: List[int] = []
    class_names: List[str] = []

    usable = min(len(horizon_returns), len(horizon_mask))
    for idx in range(usable):
        if float(horizon_mask[idx]) <= 0:
            class_targets.append(MOVEMENT_CLASS_TO_ID["flat"])
            class_mask.append(0)
            class_names.append("masked")
            continue

        class_id = classify_movement_return(
            float(horizon_returns[idx]),
            movement_label_config=movement_label_config,
        )
        class_targets.append(class_id)
        class_mask.append(1)
        class_names.append(MOVEMENT_ID_TO_CLASS[class_id])

    return {
        "movement_class_targets": class_targets,
        "movement_class_mask": class_mask,
        "movement_class_names": class_names,
    }


def compute_movement_label_config(valid_returns: Sequence[float]) -> Dict[str, float]:
    del valid_returns
    return dict(DEFAULT_MOVEMENT_LABEL_CONFIG)