from __future__ import annotations

import logging
import os
import sys
from pathlib import Path
from types import SimpleNamespace

import torch


REPO_ROOT = Path(__file__).resolve().parents[1]
ENTVIT_ROOT = Path(os.environ.get("ENTVIT_ROOT", str(REPO_ROOT.parent / "entvit")))
sys.path.insert(0, str(ENTVIT_ROOT / "scripts"))
sys.path.insert(0, str(REPO_ROOT))

import entvit_specialization_dense_eval as dense  # noqa: E402
from simdinov2.eval.setup import build_model  # noqa: E402
from simdinov2.utils.config import setup as setup_cfg  # noqa: E402


logger = logging.getLogger("simdino_entvit_dense_eval")


def build_simdino_model(args):
    cfg_args = SimpleNamespace(
        base_config=args.base_config,
        config_file=args.config_file,
        output_dir=args.output_dir,
        opts=list(args.opts or []),
        resume_from="",
        seed=args.seed,
    )
    cfg = setup_cfg(cfg_args, enable_dist=False)
    model = build_model(cfg, args.pretrained_weights)
    model.eval().cuda()
    for param in model.parameters():
        param.requires_grad_(False)
    logger.info(
        "loaded SimDINO checkpoint=%s config=%s arch=%s patch=%s blocks=%s",
        args.pretrained_weights,
        args.config_file,
        cfg.student.arch,
        model.patch_size,
        getattr(model, "n_blocks", "unknown"),
    )
    return model


def build_model_and_head(args):
    model = build_simdino_model(args)
    layer_indices = args.layer_indices or [round(((index + 1) * model.n_blocks) / 4) - 1 for index in range(4)]
    layer_indices = [min(max(int(index), 0), model.n_blocks - 1) for index in layer_indices]
    sample = torch.zeros(1, 3, args.image_size, args.image_size, device="cuda")
    feature_map, _ = dense.extract_feature_map(model, sample, args.feature_mode, layer_indices)
    in_channels = int(feature_map.shape[1])
    if args.task == "voc_seg":
        return model, dense.SegmentationLinearProbeHead(in_channels, args.num_classes).cuda(), layer_indices
    return model, dense.DepthLinearProbeHead(in_channels).cuda(), layer_indices


def get_args_parser():
    parser = dense.get_args_parser()
    parser.add_argument("--config-file", required=True)
    parser.add_argument("--base-config", default="ssl_default_config")
    parser.add_argument("--opts", nargs="*", default=[])
    return parser


def main() -> int:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    args = get_args_parser().parse_args()
    dense.build_model_and_head = build_model_and_head
    dense.seed_everything(args.seed)
    if args.task == "voc_seg" and not args.feature_mode.startswith("last_"):
        raise ValueError("VOC specialization protocol expects a final-layer feature mode.")
    if args.task == "nyuv2_depth" and not args.feature_mode.startswith("4layer_"):
        raise ValueError("NYUv2 specialization protocol expects a four-layer feature mode.")
    if args.task == "voc_seg":
        return dense.train_segmentation(args)
    return dense.train_depth(args)


if __name__ == "__main__":
    raise SystemExit(main())