#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import time
from dataclasses import asdict
from pathlib import Path
from typing import Any

import numpy as np
import torch
from transformers import LogitsProcessorList, StoppingCriteriaList
from tqdm import trange

from alpamayo_lite.action_space import RelativeDeltaActionSpace
from alpamayo_lite.config import load_config
from alpamayo_lite.data import load_manifest
from alpamayo_lite.evaluation import save_trajectory_plot, trajectory_metrics
from alpamayo_lite.expert import PrefixKV
from alpamayo_lite.prompts import build_vlm_prefix
from alpamayo_lite.runtime import (
    empty_device_cache,
    parameter_counts,
    resolve_dtype,
    seed_everything,
    select_device,
    write_json,
)
from alpamayo_lite.trajectory_tokens import (
    DiscreteTrajectoryMask,
    StopOneTokenAfterBoundary,
    TrajectoryVocabulary,
    add_trajectory_vocabulary,
    build_structured_target,
    dequantize_action,
    extract_trajectory_values,
    mask_after_boundary,
    quantize_action,
)
from alpamayo_lite.vlm import (
    configure_stage1_trainability,
    enforce_frozen_module_eval,
    load_vlm,
    prepare_stage1_batch,
    process_image_text,
)


def _resolve(root: Path, value: str) -> Path:
    path = Path(value).expanduser()
    return path.resolve() if path.is_absolute() else (root / path).resolve()


def _read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _infer_meta_action(reasoning: str) -> str:
    lowered = reasoning.lower()
    for phrase, label in (
        ("steer left", "STEER_LEFT"),
        ("turn left", "TURN_LEFT"),
        ("steer right", "STEER_RIGHT"),
        ("turn right", "TURN_RIGHT"),
        ("stop", "STOP"),
        ("slow", "SLOW_DOWN"),
    ):
        if phrase in lowered:
            return label
    return "CONTINUE"


def _vocabulary_from_tokenizer(tokenizer: Any, size: int) -> TrajectoryVocabulary:
    start_id = int(tokenizer.convert_tokens_to_ids("<i0>"))
    end_id = int(tokenizer.convert_tokens_to_ids(f"<i{size - 1}>"))
    if end_id - start_id + 1 != size:
        raise RuntimeError("Reloaded trajectory vocabulary is not contiguous")
    special_tokens = (
        "<|cot_start|>",
        "<|cot_end|>",
        "<|meta_action_start|>",
        "<|meta_action_end|>",
        "<|traj_future_start|>",
        "<|traj_future_end|>",
    )
    special_token_ids = {
        token: int(tokenizer.convert_tokens_to_ids(token)) for token in special_tokens
    }
    return TrajectoryVocabulary(size, start_id, end_id, special_token_ids)


def _cache_summary(cache: object) -> dict[str, Any]:
    prefix = PrefixKV.from_transformers_cache(cache)
    first_key = prefix.keys[0]
    sample = first_key[..., : min(8, first_key.shape[-2]), : min(8, first_key.shape[-1])]
    return {
        "layers": prefix.num_layers,
        "sequence_length": prefix.sequence_length,
        "first_key_shape": list(first_key.shape),
        "first_key_sample_sum": float(sample.float().sum().cpu()),
        "all_finite": bool(
            all(torch.isfinite(key).all() and torch.isfinite(value).all() for key, value in zip(prefix.keys, prefix.values))
        ),
    }


def _top_tokens(logits: torch.Tensor, tokenizer: Any, count: int = 5) -> list[dict[str, Any]]:
    values, ids = torch.topk(logits.float(), k=count)
    return [
        {
            "id": int(token_id),
            "token": tokenizer.convert_ids_to_tokens(int(token_id)),
            "logit": float(value),
        }
        for value, token_id in zip(values.cpu(), ids.cpu())
    ]


@torch.no_grad()
def _generate(
    model: Any,
    processor: Any,
    image: Any,
    prefix_text: str,
    device: torch.device,
    dtype: torch.dtype,
    max_new_tokens: int,
    logits_processor: LogitsProcessorList | None = None,
    stopping_criteria: StoppingCriteriaList | None = None,
) -> tuple[dict[str, Any], torch.Tensor, Any]:
    inputs = process_image_text(processor, image, prefix_text, device, dtype)
    prompt_length = inputs["input_ids"].shape[1]
    generated = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        use_cache=True,
        return_dict_in_generate=True,
        output_logits=True,
        output_scores=True,
        logits_processor=logits_processor,
        stopping_criteria=stopping_criteria,
        pad_token_id=processor.tokenizer.pad_token_id,
        eos_token_id=processor.tokenizer.eos_token_id,
    )
    new_ids = generated.sequences[:, prompt_length:]
    step_records = []
    returned_logits = generated.logits or ()
    for index, token_id in enumerate(new_ids[0].tolist()):
        if index >= len(returned_logits):
            break
        logits = returned_logits[index][0]
        log_probability = torch.log_softmax(logits.float(), dim=-1)[token_id]
        step_records.append(
            {
                "index": index,
                "id": int(token_id),
                "token": processor.tokenizer.convert_ids_to_tokens(int(token_id)),
                "selected_logit": float(logits[token_id].float().cpu()),
                "selected_log_probability": float(log_probability.cpu()),
                "top5": _top_tokens(logits, processor.tokenizer),
            }
        )
    result = {
        "prompt_length": prompt_length,
        "new_token_ids": new_ids[0].detach().cpu().tolist(),
        "raw_text": processor.tokenizer.decode(new_ids[0], skip_special_tokens=False),
        "clean_text": processor.tokenizer.decode(new_ids[0], skip_special_tokens=True),
        "steps": step_records,
        "cache": _cache_summary(generated.past_key_values),
    }
    return result, new_ids, generated


@torch.no_grad()
def _teacher_forcing_probe(
    model: Any,
    inputs: dict[str, torch.Tensor],
    labels: torch.Tensor,
    vocabulary: TrajectoryVocabulary,
    tokenizer: Any,
) -> tuple[dict[str, Any], np.ndarray, np.ndarray]:
    outputs = model(**inputs, use_cache=False, return_dict=True)
    logits = outputs.logits[:, :-1]
    next_ids = inputs["input_ids"][:, 1:]
    supervised = labels[:, 1:] != -100
    trajectory = (
        (next_ids >= vocabulary.start_id)
        & (next_ids <= vocabulary.end_id)
        & supervised
    )
    supervised_logits = logits[supervised].float()
    supervised_ids = next_ids[supervised]
    trajectory_logits = logits[trajectory].float()
    trajectory_ids = next_ids[trajectory]
    if trajectory_logits.shape[0] == 0:
        raise RuntimeError("No supervised trajectory tokens were found")

    target_logits = trajectory_logits.gather(1, trajectory_ids[:, None]).squeeze(1)
    ranks = (trajectory_logits > target_logits[:, None]).sum(dim=1) + 1
    probabilities = torch.softmax(trajectory_logits, dim=-1).gather(
        1, trajectory_ids[:, None]
    ).squeeze(1)
    top_values, top_ids = torch.topk(trajectory_logits, k=5, dim=-1)
    records = []
    for index in range(trajectory_logits.shape[0]):
        records.append(
            {
                "index": index,
                "target_id": int(trajectory_ids[index]),
                "target_token": tokenizer.convert_ids_to_tokens(int(trajectory_ids[index])),
                "target_logit": float(target_logits[index].cpu()),
                "target_probability": float(probabilities[index].cpu()),
                "target_rank": int(ranks[index].cpu()),
                "top5": [
                    {
                        "id": int(token_id),
                        "token": tokenizer.convert_ids_to_tokens(int(token_id)),
                        "logit": float(value),
                    }
                    for value, token_id in zip(top_values[index].cpu(), top_ids[index].cpu())
                ],
            }
        )

    discrete_slice = trajectory_logits[:, vocabulary.start_id : vocabulary.end_id + 1]
    result = {
        "supervised_token_count": int(supervised.sum()),
        "trajectory_supervised_token_count": int(trajectory.sum()),
        "next_token_accuracy": float(
            (supervised_logits.argmax(dim=-1) == supervised_ids).float().mean().cpu()
        ),
        "trajectory_next_token_accuracy": float(
            (trajectory_logits.argmax(dim=-1) == trajectory_ids).float().mean().cpu()
        ),
        "trajectory_mean_target_probability": float(probabilities.mean().cpu()),
        "trajectory_max_target_rank": int(ranks.max().cpu()),
        "trajectory_records": records,
    }
    return (
        result,
        discrete_slice.detach().cpu().numpy(),
        (trajectory_ids - vocabulary.start_id).detach().cpu().numpy(),
    )


@torch.no_grad()
def _handoff_probe(
    model: Any,
    processor: Any,
    image: Any,
    handoff_text: str,
    first_target_id: int,
    vocabulary: TrajectoryVocabulary,
    device: torch.device,
    dtype: torch.dtype,
) -> tuple[dict[str, Any], np.ndarray, np.ndarray]:
    inputs = process_image_text(processor, image, handoff_text, device, dtype)
    outputs = model(
        **inputs,
        use_cache=True,
        return_dict=True,
        logits_to_keep=1,
    )
    raw = outputs.logits[0, -1].float()
    masked = DiscreteTrajectoryMask(vocabulary.start_id, vocabulary.size)(
        inputs["input_ids"], raw[None].clone()
    )[0]
    target_logit = raw[first_target_id]
    target_rank = int((raw > target_logit).sum().cpu()) + 1
    outside = torch.ones_like(raw, dtype=torch.bool)
    outside[vocabulary.start_id : vocabulary.end_id + 1] = False
    result = {
        "first_target_id": first_target_id,
        "first_target_token": processor.tokenizer.convert_ids_to_tokens(first_target_id),
        "raw_first_target_logit": float(target_logit.cpu()),
        "raw_first_target_rank": target_rank,
        "raw_top5": _top_tokens(raw, processor.tokenizer),
        "processed_top5": _top_tokens(masked, processor.tokenizer),
        "raw_trajectory_finite_count": int(
            torch.isfinite(raw[vocabulary.start_id : vocabulary.end_id + 1]).sum().cpu()
        ),
        "processed_trajectory_negative_infinity_count": int(
            torch.isneginf(masked[vocabulary.start_id : vocabulary.end_id + 1]).sum().cpu()
        ),
        "outside_range_max_absolute_change": float((masked[outside] - raw[outside]).abs().max().cpu()),
        "cache": _cache_summary(outputs.past_key_values),
    }
    return (
        result,
        raw[vocabulary.start_id : vocabulary.end_id + 1].detach().cpu().numpy(),
        masked[vocabulary.start_id : vocabulary.end_id + 1].detach().cpu().numpy(),
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/local_overfit.yaml")
    parser.add_argument("--output-dir", default="outputs/stage1_vlm_trajectory")
    parser.add_argument("--steps", type=int, default=120)
    parser.add_argument("--learning-rate", type=float, default=1e-4)
    parser.add_argument("--trajectory-vocab-size", type=int, default=128)
    parser.add_argument("--action-min", type=float, default=-4.0)
    parser.add_argument("--action-max", type=float, default=4.0)
    parser.add_argument("--max-new-tokens", type=int, default=96)
    parser.add_argument("--early-stop-loss", type=float, default=1e-4)
    parser.add_argument("--early-stop-patience", type=int, default=5)
    args = parser.parse_args()
    if args.steps <= 0:
        raise ValueError("--steps must be positive")
    if args.early_stop_loss < 0:
        raise ValueError("--early-stop-loss must be non-negative")
    if args.early_stop_patience <= 0:
        raise ValueError("--early-stop-patience must be positive")

    cfg = load_config(args.config)
    seed_everything(cfg.seed)
    device = select_device()
    dtype = resolve_dtype(cfg.vlm.dtype, device)
    samples = load_manifest(cfg.resolve(cfg.paths.manifest))
    if len(samples) != 1:
        raise ValueError("This local trajectory-logit validation expects exactly one sample")
    sample = samples[0]
    history, future = sample.load_trajectory()
    image = sample.load_image()

    processor, model = load_vlm(cfg.resolve(cfg.paths.stage1_dir), cfg.vlm, device, dtype)
    tokenizer = processor.tokenizer
    original_tokenizer_length = len(tokenizer)
    vocabulary = add_trajectory_vocabulary(tokenizer, args.trajectory_vocab_size)
    model_vocab_size = model.get_output_embeddings().weight.shape[0]
    if len(tokenizer) > model_vocab_size:
        raise RuntimeError(
            f"Expanded tokenizer ({len(tokenizer)}) exceeds preallocated LM head ({model_vocab_size})"
        )
    if model.get_input_embeddings().weight.data_ptr() != model.get_output_embeddings().weight.data_ptr():
        raise RuntimeError("Trajectory validation expects tied input/output embeddings")

    expert_config = _read_json(cfg.resolve(cfg.paths.stage2_dir) / "expert_config.json")
    action_space = RelativeDeltaActionSpace(
        torch.tensor(expert_config["action_mean"], device=device),
        torch.tensor(expert_config["action_scale"], device=device),
    ).to(device)
    target_xy = future[None, :, :2].to(device)
    history_xy = history[None, :, :2].to(device)
    target_action = action_space.trajectory_to_action(history_xy, target_xy)
    quantized = quantize_action(
        target_action,
        args.trajectory_vocab_size,
        args.action_min,
        args.action_max,
    )
    meta_action = _infer_meta_action(sample.reasoning)
    structured_target = build_structured_target(sample.reasoning, meta_action, quantized)
    prefix_text = build_vlm_prefix(history)
    full_text = prefix_text + structured_target

    baseline, _, baseline_output = _generate(
        model,
        processor,
        image,
        prefix_text,
        device,
        dtype,
        args.max_new_tokens,
    )
    del baseline_output

    trainability = configure_stage1_trainability(model, cfg.stage1)
    embedding_weight = model.get_input_embeddings().weight
    embedding_weight.requires_grad_(True)
    trainable_rows = list(range(vocabulary.start_id, vocabulary.end_id + 1)) + list(
        vocabulary.special_token_ids.values()
    )
    gradient_mask = torch.zeros(
        embedding_weight.shape[0], 1, device=embedding_weight.device, dtype=embedding_weight.dtype
    )
    gradient_mask[trainable_rows] = 1
    gradient_hook = embedding_weight.register_hook(lambda gradient: gradient * gradient_mask)

    inputs, labels, prefix_length = prepare_stage1_batch(
        processor,
        image,
        prefix_text,
        full_text,
        device,
        dtype,
    )
    target_ids = inputs["input_ids"][:, prefix_length:]
    trajectory_supervision_count = int(
        (
            (target_ids >= vocabulary.start_id)
            & (target_ids <= vocabulary.end_id)
        ).sum().cpu()
    )
    if trajectory_supervision_count != quantized.numel():
        raise RuntimeError(
            f"Expected {quantized.numel()} trajectory labels, found {trajectory_supervision_count}"
        )

    optimizer = torch.optim.AdamW(
        [parameter for parameter in model.parameters() if parameter.requires_grad],
        lr=args.learning_rate,
        weight_decay=cfg.stage1.weight_decay,
        eps=cfg.stage1.adam_epsilon,
    )
    losses: list[float] = []
    accuracies: list[float] = []
    converged_steps = 0
    started = time.perf_counter()
    model.train()
    enforce_frozen_module_eval(model)
    for step in trange(args.steps, desc="vlm-trajectory-logits"):
        optimizer.zero_grad(set_to_none=True)
        outputs = model(**inputs, labels=labels, use_cache=False, return_dict=True)
        loss = outputs.loss
        if loss is None or not torch.isfinite(loss):
            raise RuntimeError(f"Non-finite trajectory-logit loss at step {step}: {loss}")
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            [parameter for parameter in model.parameters() if parameter.requires_grad], 1.0
        )
        optimizer.step()
        value = float(loss.detach().cpu())
        next_predictions = outputs.logits[:, :-1].argmax(dim=-1)
        shifted_labels = labels[:, 1:]
        supervised = shifted_labels != -100
        accuracy = float(
            (next_predictions[supervised] == shifted_labels[supervised]).float().mean().cpu()
        )
        losses.append(value)
        accuracies.append(accuracy)
        if step == 0 or (step + 1) % cfg.stage1.log_every == 0:
            print(f"trajectory-vlm step={step + 1} loss={value:.6f} accuracy={accuracy:.4f}")
        del outputs, next_predictions
        if value <= args.early_stop_loss and accuracy == 1.0:
            converged_steps += 1
        else:
            converged_steps = 0
        if converged_steps >= args.early_stop_patience:
            print(
                f"trajectory-vlm converged at step={step + 1}: "
                f"loss={value:.6f} accuracy={accuracy:.4f}"
            )
            break
    elapsed = time.perf_counter() - started

    model.gradient_checkpointing_disable()
    model.config.use_cache = True
    model.model.language_model.config.use_cache = True
    output_dir = _resolve(cfg.root, args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output_dir, safe_serialization=True, max_shard_size="2GB")
    processor.save_pretrained(output_dir)
    write_json(
        output_dir / "trajectory_vocabulary.json",
        {
            **asdict(vocabulary),
            "original_tokenizer_length": original_tokenizer_length,
            "expanded_tokenizer_length": len(tokenizer),
            "model_vocab_size": model_vocab_size,
            "preallocated_rows_used": True,
            "action_min": args.action_min,
            "action_max": args.action_max,
            "action_mean": expert_config["action_mean"],
            "action_scale": expert_config["action_scale"],
            "action_shape": expert_config["action_shape"],
        },
    )
    gradient_hook.remove()
    del optimizer, model, processor
    empty_device_cache(device)

    processor, model = load_vlm(output_dir, cfg.vlm, device, dtype)
    model.requires_grad_(False)
    model.eval()
    vocabulary = _vocabulary_from_tokenizer(processor.tokenizer, args.trajectory_vocab_size)
    inputs, labels, prefix_length = prepare_stage1_batch(
        processor,
        image,
        prefix_text,
        full_text,
        device,
        dtype,
    )
    teacher, teacher_logits, teacher_values = _teacher_forcing_probe(
        model, inputs, labels, vocabulary, processor.tokenizer
    )

    cot_and_meta = (
        f"<|cot_start|>{sample.reasoning}<|cot_end|>\n"
        f"<|meta_action_start|>{meta_action}<|meta_action_end|>\n"
    )
    handoff_text = prefix_text + cot_and_meta + "<|traj_future_start|>"
    first_target_id = vocabulary.start_id + int(quantized.reshape(-1)[0].cpu())
    handoff, boundary_raw, boundary_masked = _handoff_probe(
        model,
        processor,
        image,
        handoff_text,
        first_target_id,
        vocabulary,
        device,
        dtype,
    )

    full_generation, full_ids, full_output = _generate(
        model,
        processor,
        image,
        prefix_text,
        device,
        dtype,
        args.max_new_tokens,
    )
    generated_values = extract_trajectory_values(full_ids[0], vocabulary)
    expected_values = quantized.detach().cpu().reshape(-1)
    generated_token_accuracy = (
        float((generated_values.cpu() == expected_values).float().mean())
        if generated_values.numel() == expected_values.numel()
        else 0.0
    )
    if generated_values.numel() != target_action.numel():
        raise RuntimeError(
            f"Generated {generated_values.numel()} trajectory tokens, expected {target_action.numel()}"
        )
    generated_action = dequantize_action(
        generated_values.reshape_as(target_action).to(device),
        args.trajectory_vocab_size,
        args.action_min,
        args.action_max,
    )
    generated_xy = action_space.action_to_trajectory(generated_action, history_xy)
    quantized_target_action = dequantize_action(
        quantized,
        args.trajectory_vocab_size,
        args.action_min,
        args.action_max,
    )
    quantized_target_xy = action_space.action_to_trajectory(quantized_target_action, history_xy)
    del full_output

    boundary_id = vocabulary.special_token_ids["<|traj_future_start|>"]
    source_mask = DiscreteTrajectoryMask(vocabulary.start_id, vocabulary.size)
    source_stop = StopOneTokenAfterBoundary(boundary_id)
    truncated_generation, truncated_ids, truncated_output = _generate(
        model,
        processor,
        image,
        prefix_text,
        device,
        dtype,
        args.max_new_tokens,
        logits_processor=LogitsProcessorList([source_mask]),
        stopping_criteria=StoppingCriteriaList([source_stop]),
    )
    visible_ids = mask_after_boundary(
        truncated_output.sequences,
        boundary_id,
        processor.tokenizer.pad_token_id,
    )
    truncated_generation["visible_new_token_ids"] = visible_ids[
        :, truncated_generation["prompt_length"] :
    ][0].cpu().tolist()
    truncated_generation["visible_raw_text"] = processor.tokenizer.decode(
        visible_ids[:, truncated_generation["prompt_length"] :][0],
        skip_special_tokens=False,
    )
    truncated_generation["boundary_generated"] = bool((truncated_ids == boundary_id).any())
    truncated_generation["trajectory_tokens_visible"] = int(
        (
            (visible_ids >= vocabulary.start_id)
            & (visible_ids <= vocabulary.end_id)
        ).sum().cpu()
    )
    boundary_positions = torch.nonzero(truncated_ids[0] == boundary_id, as_tuple=False).flatten()
    post_boundary_logit_index = (
        int(boundary_positions[0]) + 1 if boundary_positions.numel() else -1
    )
    if 0 <= post_boundary_logit_index < len(truncated_output.logits or ()):
        post_boundary_logits = truncated_output.logits[post_boundary_logit_index][0]
        truncated_generation["returned_post_boundary_trajectory_negative_infinity_count"] = int(
            torch.isneginf(
                post_boundary_logits[vocabulary.start_id : vocabulary.end_id + 1]
            ).sum().cpu()
        )
    del truncated_output

    save_trajectory_plot(
        output_dir / "vlm_discrete_trajectory.png",
        history[:, :2],
        target_xy[0],
        generated_xy[0],
        prediction_label="VLM discrete tokens",
        title="VLM discrete-trajectory logits",
    )
    np.savez_compressed(
        output_dir / "trajectory_logits.npz",
        teacher_trajectory_logits=teacher_logits,
        teacher_target_values=teacher_values,
        boundary_raw_trajectory_logits=boundary_raw,
        boundary_masked_trajectory_logits=boundary_masked,
        target_action=target_action.detach().cpu().numpy(),
        quantized_target_values=quantized.detach().cpu().numpy(),
        generated_values=generated_values.detach().cpu().numpy(),
        generated_action=generated_action.detach().cpu().numpy(),
        target_xy=target_xy.detach().cpu().numpy(),
        generated_xy=generated_xy.detach().cpu().numpy(),
    )

    metrics = {
        "stage": "vlm_discrete_trajectory_logits",
        "device": str(device),
        "dtype": str(dtype),
        "model_parameters": parameter_counts(model),
        "vocabulary": asdict(vocabulary),
        "original_tokenizer_length": original_tokenizer_length,
        "expanded_tokenizer_length": len(processor.tokenizer),
        "model_vocab_size": model_vocab_size,
        "semantic_trajectory_supervision_count": trajectory_supervision_count,
        "structured_target": structured_target,
        "target_meta_action": meta_action,
        "training": {
            "steps": len(losses),
            "requested_steps": args.steps,
            "learning_rate": args.learning_rate,
            "elapsed_seconds": elapsed,
            "initial_loss": losses[0],
            "final_loss": losses[-1],
            "minimum_loss": min(losses),
            "final_training_accuracy": accuracies[-1],
            "losses": losses,
            "accuracies": accuracies,
            "base_trainability": trainability,
            "trainable_token_rows": len(trainable_rows),
        },
        "baseline_generation_before_structured_training": baseline,
        "teacher_forcing": teacher,
        "raw_handoff_logits_and_cache": handoff,
        "untruncated_vlm_generation": {
            **full_generation,
            "trajectory_token_count": int(generated_values.numel()),
            "trajectory_token_accuracy": generated_token_accuracy,
            "generated_values": generated_values.cpu().tolist(),
            "target_values": expected_values.tolist(),
            "action_mse": float(
                torch.nn.functional.mse_loss(generated_action, target_action).cpu()
            ),
            **trajectory_metrics(generated_xy, target_xy),
        },
        "quantization_only": trajectory_metrics(quantized_target_xy, target_xy),
        "source_style_diffusion_handoff_generation": truncated_generation,
    }
    write_json(output_dir / "trajectory_logits_metrics.json", metrics)
    print(json.dumps(metrics, indent=2, ensure_ascii=False))
    del model, processor
    empty_device_cache(device)


if __name__ == "__main__":
    main()