import matplotlib.pyplot as plt
import numpy as np
from typing import List
from ..utils.dataclass import TTLBlock
from ..hardware.constraints import HardwareConstraints


def plot_ttl_timeline(
        ttl_blocks: List[TTLBlock],
        original_blocks: List[dict],
        tick_duration: float = 20e-9
):
    """
    Визуализирует TTL timeline + оригинальные сигналы:
    RF и ADC отображаются как синусоиды, GRAD как трапеция.
    Оригинальные сигналы сдвигаются вправо с учётом начального TTL-блока.

    :param ttl_blocks: список TTLBlock объектов (результат Synchronizer)
    :param original_blocks: список исходных блоков {"type": [...], "duration": float}
    :param tick_duration: длительность одного тика (сек)
    """
    # === TTL timeline ===
    time = 0
    ttl_stream = {"RF": [], "SW": [], "ADC": [], "GR": []}
    time_stream = []

    for block in ttl_blocks:
        for _ in range(block.duration):
            time_stream.append(time)
            ttl_stream["RF"].append(block.RF)
            ttl_stream["SW"].append(block.SW)
            ttl_stream["ADC"].append(block.ADC)
            ttl_stream["GR"].append(block.GR)
            time += tick_duration

    # === Исходные сигналы (RF, ADC — синус; GRAD — трапеция) ===
    signal_plot = {"RF": [], "ADC": [], "GRAD": []}
    signal_time = []
    t = 0

    for block in original_blocks:
        duration = block["duration"]
        ticks = int(duration / tick_duration)
        t_block = np.linspace(t, t + duration, ticks)

        # RF синус
        if "RF" in block["type"]:
            signal_plot["RF"].extend(np.sin(2 * np.pi * 20e6 * t_block))  # 20 кГц
        else:
            signal_plot["RF"].extend([0] * ticks)

        # ADC синус, меньшей амплитудой
        if "ADC" in block["type"]:
            signal_plot["ADC"].extend(0.5 * np.sin(2 * np.pi * 10e6 * t_block))
        else:
            signal_plot["ADC"].extend([0] * ticks)

        # GRAD — трапеция
        if "GRAD" in block["type"]:
            ramp_len = ticks // 4
            plateau_len = ticks // 2
            trapezoid = list(np.linspace(0, 1, ramp_len)) + \
                        [1] * plateau_len + \
                        list(np.linspace(1, 0, ticks - ramp_len - plateau_len))
            signal_plot["GRAD"].extend(trapezoid)
        else:
            signal_plot["GRAD"].extend([0] * ticks)

        signal_time.extend(t_block)
        t += duration
    hw = HardwareConstraints()
    # === Сдвиг исходных сигналов на START_DELAY ===
    start_offset = tick_duration * (
            round(hw.START_DELAY / tick_duration) +
            max(
                round(hw.RF_DELAY / tick_duration),
                round(hw.TR_DELAY / tick_duration),
                round(hw.GRAD_DELAY / tick_duration)
            )
    )
    adjusted_signal_time = [t + start_offset for t in signal_time]

    # === Визуализация ===
    plt.figure(figsize=(14, 8))

    # TTL сигналы — логика
    for i, (label, values) in enumerate(ttl_stream.items()):
        plt.step(
            [t * 1e3 for t in time_stream],
            [v + i * 2 for v in values],
            where='post',
            label=f"TTL: {label}"
        )

    # Исходные сигналы (синус/трапеция)
    offset_map = {"RF": 8, "ADC": 10, "GRAD": 12}
    color_map = {"RF": "blue", "ADC": "purple", "GRAD": "orange"}

    for label in signal_plot:
        y = np.array(signal_plot[label]) + offset_map[label]
        plt.plot(
            [t * 1e3 for t in adjusted_signal_time],
            y,
            label=f"Signal: {label}",
            color=color_map[label]
        )

    # Финальные штрихи
    plt.title("TTL Timeline + Original Signal Shapes (с учётом задержек)")
    plt.xlabel("Time (mks)")
    yticks = [i * 2 for i in range(4)] + list(offset_map.values())
    ylabels = list(ttl_stream.keys()) + list(offset_map.keys())
    plt.yticks(yticks, ylabels)
    plt.grid(True)
    plt.legend(loc="upper right")
    plt.tight_layout()
    plt.show()