from __future__ import annotations

import math
import sys
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from comfyui_spectrum.config import SpectrumConfig
from comfyui_spectrum.forecast import ChebyshevSpectrumForecaster
from comfyui_spectrum.runtime import SpectrumRuntime


def make_runtime() -> SpectrumRuntime:
    cfg = SpectrumConfig(
        blend_weight=0.5,
        degree=4,
        ridge_lambda=0.1,
        window_size=2.0,
        flex_window=0.75,
        warmup_steps=5,
        max_history=128,
    ).validate()
    return SpectrumRuntime(cfg)


def test_solver_step_scheduler() -> None:
    runtime = make_runtime()
    sample_sigmas = torch.linspace(1.0, 0.0, 51)
    run_id = runtime.start_run(sample_sigmas, "sample_euler", supports_solver_steps=True)
    total_steps = len(sample_sigmas) - 1

    for step_id in range(5):
        decision = runtime.begin_solver_step(
            run_id,
            step_id,
            runtime.time_coord_for_step(step_id),
            total_steps,
        )
        assert decision["actual_forward"] is True
        runtime.register_model_hook_call(
            run_id,
            step_id,
            expected_shape=(1, 8, 4),
            branch_signature=(("cond_or_uncond", (0, 1)),),
        )
        runtime.observe_actual_feature(decision["run_id"], decision["solver_step_id"], torch.randn(1, 8, 4))
        runtime.finalize_solver_step(decision["run_id"], decision["solver_step_id"], used_forecast=False)

    decision = runtime.begin_solver_step(
        run_id,
        5,
        runtime.time_coord_for_step(5),
        total_steps,
    )
    runtime.register_model_hook_call(
        run_id,
        5,
        expected_shape=(1, 8, 4),
        branch_signature=(("cond_or_uncond", (0, 1)),),
    )
    predicted = runtime.predict_feature(run_id, 5, expected_shape=(1, 8, 4))
    assert predicted is not None
    assert predicted.shape == (1, 8, 4)
    runtime.finalize_solver_step(run_id, 5, used_forecast=not decision["actual_forward"])
    runtime.end_run(run_id)


def test_forecast_fallback_reconciles_bookkeeping() -> None:
    runtime = make_runtime()
    sample_sigmas = torch.linspace(1.0, 0.0, 51)
    run_id = runtime.start_run(sample_sigmas, "sample_euler", supports_solver_steps=True)
    total_steps = len(sample_sigmas) - 1

    for step_id in range(5):
        decision = runtime.begin_solver_step(
            run_id,
            step_id,
            runtime.time_coord_for_step(step_id),
            total_steps,
        )
        runtime.register_model_hook_call(run_id, step_id, expected_shape=(1, 8, 4))
        runtime.observe_actual_feature(run_id, step_id, torch.randn(1, 8, 4))
        runtime.finalize_solver_step(run_id, step_id, used_forecast=False)

    decision = runtime.begin_solver_step(
        run_id,
        5,
        runtime.time_coord_for_step(5),
        total_steps,
    )
    assert decision["actual_forward"] is False
    before_window = runtime.curr_ws

    runtime.register_model_hook_call(run_id, 5, expected_shape=(1, 8, 4))
    predicted = runtime.predict_feature(run_id, 5, expected_shape=(2, 8, 4))
    assert predicted is None
    runtime.observe_actual_feature(run_id, 5, torch.randn(1, 8, 4))
    runtime.finalize_solver_step(run_id, 5, used_forecast=False)

    assert runtime.stats.forecasted_count == 0
    assert runtime.stats.actual_forward_count == 6
    assert runtime.num_consecutive_cached_steps == 0
    assert decision["actual_forward"] is True
    assert runtime.curr_ws == before_window
    runtime.end_run(run_id)


def test_observe_actual_feature_clears_forecast_latch() -> None:
    runtime = make_runtime()
    sample_sigmas = torch.linspace(1.0, 0.0, 51)
    run_id = runtime.start_run(sample_sigmas, "sample_euler", supports_solver_steps=True)
    total_steps = len(sample_sigmas) - 1

    for step_id in range(5):
        decision = runtime.begin_solver_step(
            run_id,
            step_id,
            runtime.time_coord_for_step(step_id),
            total_steps,
        )
        runtime.register_model_hook_call(run_id, step_id, expected_shape=(1, 8, 4))
        runtime.observe_actual_feature(run_id, step_id, torch.randn(1, 8, 4))
        runtime.finalize_solver_step(decision["run_id"], decision["solver_step_id"], used_forecast=False)

    decision = runtime.begin_solver_step(
        run_id,
        5,
        runtime.time_coord_for_step(5),
        total_steps,
    )
    runtime.register_model_hook_call(run_id, 5, expected_shape=(1, 8, 4))
    predicted = runtime.predict_feature(run_id, 5, expected_shape=(1, 8, 4))
    assert predicted is not None
    assert runtime.step_used_forecast(run_id, 5) is True
    runtime.observe_actual_feature(run_id, 5, torch.randn(1, 8, 4))
    assert runtime.step_used_forecast(run_id, 5) is False
    runtime.finalize_solver_step(decision["run_id"], decision["solver_step_id"], used_forecast=False)
    runtime.end_run(run_id)


def test_unsupported_sampler_disables_forecast() -> None:
    runtime = make_runtime()
    sample_sigmas = torch.linspace(1.0, 0.0, 51)
    run_id = runtime.start_run(sample_sigmas, "sample_heun", supports_solver_steps=False)
    decision = runtime.begin_solver_step(
        run_id,
        0,
        runtime.time_coord_for_step(0),
        len(sample_sigmas) - 1,
    )
    assert decision["actual_forward"] is True
    assert runtime.stats.forecast_disabled is True
    assert runtime.stats.disable_reason == "sampler 'sample_heun' does not expose one predict_noise call per solver step"
    runtime.register_model_hook_call(run_id, 0, expected_shape=(1, 8, 4))
    runtime.observe_actual_feature(run_id, 0, torch.randn(1, 8, 4))
    runtime.finalize_solver_step(run_id, 0, used_forecast=False)
    runtime.end_run(run_id)


def test_inconsistent_hook_shape_disables_forecast() -> None:
    runtime = make_runtime()
    sample_sigmas = torch.linspace(1.0, 0.0, 51)
    run_id = runtime.start_run(sample_sigmas, "sample_euler", supports_solver_steps=True)
    total_steps = len(sample_sigmas) - 1

    for step_id in range(5):
        decision = runtime.begin_solver_step(
            run_id,
            step_id,
            runtime.time_coord_for_step(step_id),
            total_steps,
        )
        runtime.register_model_hook_call(run_id, step_id, expected_shape=(1, 8, 4))
        runtime.observe_actual_feature(run_id, step_id, torch.randn(1, 8, 4))
        runtime.finalize_solver_step(run_id, step_id, used_forecast=False)

    decision = runtime.begin_solver_step(
        run_id,
        5,
        runtime.time_coord_for_step(5),
        total_steps,
    )
    runtime.register_model_hook_call(run_id, 5, expected_shape=(1, 8, 4))
    predicted = runtime.predict_feature(run_id, 5, expected_shape=(2, 8, 4))
    assert predicted is None
    assert runtime.stats.forecast_disabled is True
    assert runtime.stats.disable_reason == "predicted feature shape did not match the current solver-step input"
    runtime.observe_actual_feature(run_id, 5, torch.randn(1, 8, 4))
    runtime.finalize_solver_step(decision["run_id"], decision["solver_step_id"], used_forecast=False)
    runtime.end_run(run_id)


def test_multiple_hook_calls_disable_forecast() -> None:
    runtime = make_runtime()
    sample_sigmas = torch.linspace(1.0, 0.0, 51)
    run_id = runtime.start_run(sample_sigmas, "sample_euler", supports_solver_steps=True)
    total_steps = len(sample_sigmas) - 1

    for step_id in range(5):
        decision = runtime.begin_solver_step(
            run_id,
            step_id,
            runtime.time_coord_for_step(step_id),
            total_steps,
        )
        runtime.register_model_hook_call(run_id, step_id, expected_shape=(1, 8, 4))
        runtime.observe_actual_feature(run_id, step_id, torch.randn(1, 8, 4))
        runtime.finalize_solver_step(run_id, step_id, used_forecast=False)

    decision = runtime.begin_solver_step(
        run_id,
        5,
        runtime.time_coord_for_step(5),
        total_steps,
    )
    runtime.register_model_hook_call(run_id, 5, expected_shape=(1, 8, 4))
    runtime.register_model_hook_call(run_id, 5, expected_shape=(1, 8, 4))
    assert runtime.stats.forecast_disabled is True
    assert runtime.stats.disable_reason == "multiple model-hook calls observed within one solver step"
    runtime.observe_actual_feature(run_id, 5, torch.randn(1, 8, 4))
    runtime.finalize_solver_step(decision["run_id"], decision["solver_step_id"], used_forecast=False)
    runtime.end_run(run_id)


def test_nonuniform_schedule_coords_are_used() -> None:
    runtime = make_runtime()
    sample_sigmas = torch.tensor([10.0, 9.0, 1.0, 0.0])
    run_id = runtime.start_run(sample_sigmas, "sample_euler", supports_solver_steps=True)

    coords = [runtime.time_coord_for_step(i) for i in range(3)]
    assert math.isclose(coords[0], -1.0)
    assert math.isclose(coords[2], 1.0)
    assert math.isclose(coords[1], -7.0 / 9.0)
    assert not math.isclose(coords[1] - coords[0], coords[2] - coords[1])

    runtime.end_run(run_id)


def test_forecaster_respects_nonuniform_coords() -> None:
    forecaster = ChebyshevSpectrumForecaster(degree=1, ridge_lambda=0.1, max_history=16)
    forecaster.update(-1.0, torch.tensor([10.0]))
    forecaster.update(-7.0 / 9.0, torch.tensor([9.0]))

    pred = forecaster.predict(time_coord=1.0, blend_weight=0.0)
    assert torch.allclose(pred, torch.tensor([1.0]), atol=1e-5)


def main() -> None:
    test_solver_step_scheduler()
    test_forecast_fallback_reconciles_bookkeeping()
    test_observe_actual_feature_clears_forecast_latch()
    test_unsupported_sampler_disables_forecast()
    test_inconsistent_hook_shape_disables_forecast()
    test_multiple_hook_calls_disable_forecast()
    test_nonuniform_schedule_coords_are_used()
    test_forecaster_respects_nonuniform_coords()
    print("ok")


if __name__ == "__main__":
    main()