from __future__ import annotations

import json
import sys
import tempfile
import unittest
from pathlib import Path
from unittest import mock

import numpy as np
import open3d as o3d

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from train_multiview_point_transformer import (  # noqa: E402
    MultiviewLocalPointTransformer,
    load_fusion_features,
    main,
    sha256,
    spatial_split,
)


class MultiviewPointTransformerTests(unittest.TestCase):
    def make_fixture(self, root: Path) -> Path:
        artifact = root / "multiview-feature-fixture"
        artifact.mkdir()
        # Blocks 0/6/3 deterministically map to train/validation/test.
        xyz, labels = [], []
        for code, offset in ((5, 0.0), (17, 0.2)):
            for block in (0, 3, 6):
                for sample in range(200):
                    xyz.append([block * 100.0 + 10.0 + (sample % 10) * 0.01 + offset, (sample // 10) * 0.01, float(code) / 16.0])
                    labels.append(code)
        xyz_array = np.asarray(xyz, dtype=np.float64)
        rgb = np.column_stack((np.asarray(labels) == 5, np.asarray(labels) == 17, np.full(len(labels), 0.3))).astype(np.float32)
        # This unlabeled extent point fixes the eight XY blocks at 100-unit width.
        xyz_array = np.vstack((xyz_array, [[810.0, 0.0, 0.0]]))
        rgb = np.vstack((rgb, [[0.2, 0.2, 0.2]])).astype(np.float32)
        cloud = o3d.geometry.PointCloud()
        cloud.points = o3d.utility.Vector3dVector(xyz_array)
        cloud.colors = o3d.utility.Vector3dVector(rgb)
        source = artifact / "multiview-annotation-source.ply"
        self.assertTrue(o3d.io.write_point_cloud(str(source), cloud, write_ascii=False))
        feature_mean = np.column_stack((rgb, np.zeros((len(xyz_array), 5), dtype=np.float32)))
        feature_stddev = np.full((len(xyz_array), 8), 0.05, dtype=np.float32)
        dataset = artifact / "multiview-point-features.npz"
        np.savez_compressed(dataset, xyz=xyz_array, las_rgb=rgb, photo_feature_mean=feature_mean, photo_feature_stddev=feature_stddev, visible_view_count=np.full(len(xyz_array), 12, dtype=np.uint16), feature_names=np.asarray(("photo_r", "photo_g", "photo_b", "hue_sin", "hue_cos", "saturation", "gradient_magnitude", "local_intensity_stddev")))
        (artifact / "run_metadata.json").write_text(json.dumps({"annotation_source": {"kind": "multiview_photo_feature_fusion", "point_cloud": source.name, "feature_dataset": dataset.name, "point_count": len(xyz_array), "point_cloud_sha256": sha256(source), "feature_dataset_sha256": sha256(dataset)}}), encoding="utf-8")
        annotation = root / "annotation.json"
        annotation.write_text(json.dumps({"schema_version": 1, "source_path": str(source), "source_sha256": sha256(source), "labels": [[index, code] for index, code in enumerate(labels)], "class_schema": {"5": {"code": 5, "key": "vegetation", "label": "植被", "color": [59, 163, 87]}, "17": {"code": 17, "key": "transformer", "label": "变压器", "color": [30, 144, 255]}}}), encoding="utf-8")
        return annotation

    def test_loads_verified_features_and_local_attention_contract(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            annotation = self.make_fixture(Path(temp_dir))
            xyz, features, contract = load_fusion_features(annotation)
            self.assertEqual(xyz.shape, (1201, 3))
            self.assertEqual(features.shape, (1201, 23))
            self.assertEqual(contract["point_count"], 1201)
            train, validation, test = spatial_split(xyz, np.arange(len(xyz)))
            self.assertTrue(len(train) and len(validation) and len(test))
            model = MultiviewLocalPointTransformer(features.shape[1], 2)
            logits = model(__import__("torch").from_numpy(features[:4, None, :].repeat(4, axis=1)), __import__("torch").zeros((4, 4, 3)))
            self.assertEqual(tuple(logits.shape), (4, 2))

    def test_cli_writes_multiview_model_and_metrics(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            annotation = self.make_fixture(root)
            output = root / "output"
            with mock.patch.object(sys, "argv", ["trainer", "--annotation", str(annotation), "--output", str(output), "--device", "cpu", "--epochs", "1", "--batch-size", "256", "--neighbours", "4"]):
                self.assertEqual(main(), 0)
            metrics = json.loads((output / "metrics.json").read_text(encoding="utf-8"))
            self.assertTrue((output / "model.pt").is_file())
            self.assertTrue((output / "predicted-semantic-preview.ply").is_file())
            self.assertEqual(metrics["model_input_kind"], "multiview_photo_feature_fusion")
            self.assertEqual(metrics["classes"]["17"]["key"], "transformer")
            self.assertIn("not official Point Transformer V3", metrics["model"])

    def test_rejects_changed_feature_dataset(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            annotation = self.make_fixture(Path(temp_dir))
            dataset = annotation.parent / "multiview-feature-fixture" / "multiview-point-features.npz"
            dataset.write_bytes(b"changed")
            with self.assertRaisesRegex(ValueError, "checksum changed"):
                load_fusion_features(annotation)


if __name__ == "__main__":
    unittest.main()