import json
import random
import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")

import matplotlib.pyplot as plt
from PIL import Image


def visualize_crops(root, n=10, seed=42, output="crop_grid.png"):
    root = Path(root)

    # Find existing satellite, mask, and GeoJSON triplets
    triplets = []

    for sat in sorted(root.rglob("*_sat.png")):
        prefix = sat.name.removesuffix("_sat.png")

        gt = sat.with_name(f"{prefix}_gt.png")
        geo = sat.with_name(f"{prefix}_roads.geojson")

        if gt.exists() and geo.exists():
            triplets.append((sat, gt, geo))

    print(f"Found {len(triplets)} matching crops")

    if not triplets:
        raise RuntimeError("No matching crops found!")

    # Randomly select existing crops
    selected = random.Random(seed).sample(
        triplets, min(n, len(triplets))
    )

    fig, axes = plt.subplots(
        len(selected),
        3,
        figsize=(15, 4 * len(selected)),
        squeeze=False
    )

    for i, (sat_path, gt_path, geo_path) in enumerate(selected):

        sat = Image.open(sat_path).convert("RGB")
        gt = Image.open(gt_path).convert("L")

        with open(geo_path) as f:
            geojson = json.load(f)

        # Column 1: Satellite image
        axes[i, 0].imshow(sat)
        axes[i, 0].set_title(
            f"{sat_path.stem}\nSatellite",
            fontsize=9
        )

        # Column 2: Ground truth mask
        axes[i, 1].imshow(
            gt,
            cmap="gray",
            vmin=0,
            vmax=255
        )
        axes[i, 1].set_title("GT Mask")

        # Column 3: GeoJSON overlay on satellite
        ax = axes[i, 2]
        ax.imshow(sat)

        edge_count = 0

        for feature in geojson.get("features", []):
            geometry = feature.get("geometry")

            if not geometry:
                continue

            if geometry["type"] == "LineString":
                lines = [geometry["coordinates"]]

            elif geometry["type"] == "MultiLineString":
                lines = geometry["coordinates"]

            else:
                continue

            for coords in lines:
                if len(coords) < 2:
                    continue

                xs = [point[0] for point in coords]
                ys = [point[1] for point in coords]

                ax.plot(
                    xs,
                    ys,
                    color="red",
                    linewidth=1.5
                )

                edge_count += 1

        ax.set_xlim(-0.5, sat.width - 0.5)
        ax.set_ylim(sat.height - 0.5, -0.5)

        ax.set_title(
            f"GeoJSON Overlay ({edge_count} lines)"
        )

        # Remove axes
        for a in axes[i]:
            a.axis("off")

    plt.tight_layout()

    plt.savefig(
        output,
        dpi=150,
        bbox_inches="tight"
    )

    plt.close()

    print(f"Saved: {output}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--root",
        type=str,
        default="crops/val",
        help="Directory containing processed crops"
    )

    parser.add_argument(
        "--n",
        type=int,
        default=10,
        help="Number of crops to visualize"
    )

    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed"
    )

    parser.add_argument(
        "--output",
        type=str,
        default="crop_grid.png",
        help="Output visualization filename"
    )

    args = parser.parse_args()

    visualize_crops(
        root=args.root,
        n=args.n,
        seed=args.seed,
        output=args.output
    )