"""Comprehensive tests for ForestNetNonGeo dataset to maximize code coverage."""

import json
import tempfile
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock
import datetime

import pytest
import torch
import numpy as np
import pandas as pd
from PIL import Image
import albumentations as A

from terratorch.datasets.forestnet import ForestNetNonGeo


@pytest.fixture
def temp_forestnet_dataset():
    """Create a temporary ForestNet-like dataset structure."""
    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)
        
        # Create CSV files for train/val/test splits
        for split in ["train", "test", "val"]:
            csv_data = pd.DataFrame({
                "example_path": [f"event_{i:03d}" for i in range(10)],
                "merged_label": ["Plantation", "Smallholder agriculture", "Grassland shrubland", "Other"] * 2 + ["Plantation", "Other"]
            })
            csv_file = tmpdir / f"{split}_filtered.csv"
            csv_data.to_csv(csv_file, index=False)
        
        # Create event directories with images and metadata
        for i in range(10):
            event_dir = tmpdir / f"event_{i:03d}"
            event_dir.mkdir(parents=True)
            
            # Create images directories
            visible_dir = event_dir / "images" / "visible"
            infrared_dir = event_dir / "images" / "infrared"
            visible_dir.mkdir(parents=True)
            infrared_dir.mkdir(parents=True)
            
            # Create auxiliary directory with OSM data
            aux_dir = event_dir / "auxiliary"
            aux_dir.mkdir(parents=True)
            osm_file = aux_dir / "osm.json"
            osm_data = {
                "closest_city": {
                    "lat": f"{40 + i * 0.1}",
                    "lon": f"{-100 + i * 0.1}"
                }
            }
            with open(osm_file, 'w') as f:
                json.dump(osm_data, f)
            
            # Create visible and infrared images with different cloud coverage and years
            for year in [2018, 2019, 2020]:
                for cloud in [10, 30, 50]:
                    # Visible images (PNG)
                    vis_file = visible_dir / f"{year}_06_15_cloud_{cloud}.png"
                    img = Image.new('RGB', (64, 64), color=(128, 128, 128))
                    img.save(vis_file)
                    
                    # Infrared images (NPY)
                    inf_file = infrared_dir / f"{year}_06_15_cloud_{cloud}.npy"
                    inf_data = np.random.rand(64, 64, 3).astype(np.float32)
                    np.save(inf_file, inf_data)
        
        yield str(tmpdir)


class TestForestNetNonGeoInitialization:
    """Test initialization and configuration."""
    
    def test_basic_initialization(self, temp_forestnet_dataset):
        """Test basic initialization with default parameters."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        assert len(dataset) == 10
        assert dataset.split == "train"
        assert len(dataset.bands) == 6  # All bands by default
        assert dataset.use_metadata == False
    
    def test_invalid_split(self, temp_forestnet_dataset):
        """Test that invalid split raises ValueError."""
        with pytest.raises(ValueError, match="Incorrect split"):
            ForestNetNonGeo(data_root=temp_forestnet_dataset, split="invalid")
    
    def test_custom_bands(self, temp_forestnet_dataset):
        """Test initialization with custom bands."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            bands=("RED", "GREEN", "BLUE")
        )
        
        assert len(dataset.bands) == 3
        assert dataset.band_indices == [0, 1, 2]
    
    def test_invalid_bands(self, temp_forestnet_dataset):
        """Test that invalid bands raise ValueError."""
        with pytest.raises(ValueError):
            ForestNetNonGeo(
                data_root=temp_forestnet_dataset,
                split="train",
                bands=("INVALID_BAND",)
            )
    
    def test_val_split(self, temp_forestnet_dataset):
        """Test validation split."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="val")
        assert dataset.split == "val"
        assert len(dataset) == 10
    
    def test_test_split(self, temp_forestnet_dataset):
        """Test test split."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="test")
        assert dataset.split == "test"
        assert len(dataset) == 10
    
    def test_fraction_parameter(self, temp_forestnet_dataset):
        """Test dataset fraction parameter."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            fraction=0.5
        )
        
        # With stratified sampling at 50%, should have ~5 samples
        assert len(dataset) == 5
    
    def test_full_dataset_fraction(self, temp_forestnet_dataset):
        """Test that fraction=1.0 uses full dataset."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            fraction=1.0
        )
        
        assert len(dataset) == 10
    
    def test_custom_label_map(self, temp_forestnet_dataset):
        """Test custom label mapping."""
        custom_map = {
            "Plantation": 10,
            "Smallholder agriculture": 20,
            "Grassland shrubland": 30,
            "Other": 40,
        }
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            label_map=custom_map
        )
        
        assert dataset.label_map == custom_map
    
    def test_use_metadata_flag(self, temp_forestnet_dataset):
        """Test use_metadata flag."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            use_metadata=True
        )
        
        assert dataset.use_metadata == True
    
    def test_custom_transform(self, temp_forestnet_dataset):
        """Test custom transform."""
        transform = A.Compose([A.HorizontalFlip(p=1.0)])
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            transform=transform
        )
        
        assert dataset.transform == transform


class TestForestNetNonGeoGetItem:
    """Test __getitem__ and data loading."""
    
    def test_basic_getitem(self, temp_forestnet_dataset):
        """Test basic sample retrieval."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        sample = dataset[0]
        
        assert "image" in sample
        assert "label" in sample
        # After default_transform, image is torch.Tensor with shape (C, T, H, W)
        assert isinstance(sample["image"], torch.Tensor)
        assert isinstance(sample["label"], (int, torch.Tensor))
    
    def test_getitem_with_metadata(self, temp_forestnet_dataset):
        """Test sample retrieval with metadata."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            use_metadata=True
        )
        sample = dataset[0]
        
        assert "image" in sample
        assert "label" in sample
        assert "location_coords" in sample
        assert "temporal_coords" in sample
        assert isinstance(sample["location_coords"], torch.Tensor)
        assert isinstance(sample["temporal_coords"], torch.Tensor)
    
    def test_image_shape(self, temp_forestnet_dataset):
        """Test that images have correct shape."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        sample = dataset[0]
        
        # After default_transform: (C, T, H, W)
        image = sample["image"]
        assert image.ndim == 4  # (C, T, H, W)
        assert image.shape[0] == 6  # 6 bands by default
        assert image.shape[1] == 3  # 3 timesteps
    
    def test_custom_bands_filtering(self, temp_forestnet_dataset):
        """Test that custom bands filter correctly."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            bands=("RED", "GREEN", "BLUE")
        )
        sample = dataset[0]
        
        # Should only have 3 bands (RGB) - shape is (C, T, H, W)
        assert sample["image"].shape[0] == 3
    
    def test_label_mapping(self, temp_forestnet_dataset):
        """Test label mapping works correctly."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        # Test all samples
        for i in range(len(dataset)):
            sample = dataset[i]
            label = sample["label"]
            # Label should be one of the mapped values
            assert label in [0, 1, 2, 3]
    
    def test_with_transform(self, temp_forestnet_dataset):
        """Test that transforms are applied."""
        transform = A.Compose([
            A.HorizontalFlip(p=1.0)  # Always flip
        ])
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            transform=transform
        )
        sample = dataset[0]
        
        assert "image" in sample
        # Image should be transformed (though hard to verify flip without reference)


class TestForestNetNonGeoImageSelection:
    """Test image selection logic."""
    
    def test_least_cloudy_image(self, temp_forestnet_dataset):
        """Test that least cloudy images are selected."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        # Create fake image files with different cloud values
        image_files = [
            f"/path/2018_06_15_cloud_50.png",
            f"/path/2018_06_15_cloud_10.png",  # Should be selected (lowest)
            f"/path/2018_06_15_cloud_30.png",
            f"/path/2019_06_15_cloud_20.png",
            f"/path/2019_06_15_cloud_40.png",
        ]
        
        selected = dataset.least_cloudy_image(image_files)
        
        # Should select one per year
        assert len(selected) == 2  # 2018 and 2019
        assert "/path/2018_06_15_cloud_10.png" in selected
        assert "/path/2019_06_15_cloud_20.png" in selected
    
    def test_match_timesteps_sufficient_images(self, temp_forestnet_dataset):
        """Test timestep matching with sufficient images."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        image_files = [f"/path/img{i}.png" for i in range(10)]
        selected_images = [f"/path/img{i}.png" for i in range(5)]
        
        matched = dataset.match_timesteps(image_files, selected_images)
        
        # Should return exactly 3 images
        assert len(matched) == 3
    
    def test_match_timesteps_insufficient_images(self, temp_forestnet_dataset):
        """Test timestep matching with insufficient images."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        image_files = [f"/path/img{i}.png" for i in range(5)]
        selected_images = [f"/path/img0.png"]
        
        matched = dataset.match_timesteps(image_files, selected_images)
        
        # Should pad to 3 images
        assert len(matched) == 3
    
    def test_match_timesteps_no_images(self, temp_forestnet_dataset):
        """Test timestep matching with no selected images."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        image_files = [f"/path/img{i}.png" for i in range(5)]
        selected_images = []
        
        matched = dataset.match_timesteps(image_files, selected_images)
        
        # Should have 3 images from available files
        assert len(matched) == 3
    
    def test_select_images_integration(self, temp_forestnet_dataset):
        """Test full image selection pipeline."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        # Create realistic image file paths
        image_files = [
            f"/path/2018_06_15_cloud_{cloud}.png"
            for cloud in [10, 30, 50]
        ] + [
            f"/path/2019_06_15_cloud_{cloud}.png"
            for cloud in [20, 40, 60]
        ]
        
        selected = dataset.select_images(image_files)
        
        # Should return exactly 3 images
        assert len(selected) == 3


class TestForestNetNonGeoHelperMethods:
    """Test helper methods."""
    
    def test_get_coords(self, temp_forestnet_dataset):
        """Test coordinate extraction from OSM data."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        event_path = Path(temp_forestnet_dataset) / "event_000"
        coords = dataset._get_coords(event_path)
        
        assert isinstance(coords, torch.Tensor)
        assert coords.shape == (2,)  # [lat, lon]
        assert coords.dtype == torch.float32
    
    def test_get_dates(self, temp_forestnet_dataset):
        """Test date extraction from filenames."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        image_files = [
            "/path/2018_06_15_cloud_10.png",
            "/path/2019_03_20_cloud_20.png",
            "/path/2020_12_01_cloud_30.png",
        ]
        
        dates = dataset._get_dates(image_files)
        
        assert isinstance(dates, torch.Tensor)
        assert dates.shape == (3, 2)  # 3 dates, each with [year, julian_day]
        assert dates[0, 0] == 2018
        assert dates[1, 0] == 2019
        assert dates[2, 0] == 2020
    
    def test_map_label(self, temp_forestnet_dataset):
        """Test label mapping method."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        # Test mapping for first sample (should be "Plantation" = 0)
        label = dataset.map_label(0)
        assert label == 0
        
        # Test mapping for second sample (should be "Smallholder agriculture" = 1)
        label = dataset.map_label(1)
        assert label == 1
    
    def test_load_images(self, temp_forestnet_dataset):
        """Test image loading method."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            use_metadata=True
        )
        
        path = Path(temp_forestnet_dataset) / "event_000"
        vis_images, inf_images, dates = dataset._load_images(path)
        
        assert len(vis_images) == 3  # 3 timesteps
        assert len(inf_images) == 3
        assert dates is not None
        assert dates.shape[0] == 3
    
    def test_load_images_without_metadata(self, temp_forestnet_dataset):
        """Test image loading without metadata."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            use_metadata=False
        )
        
        path = Path(temp_forestnet_dataset) / "event_000"
        vis_images, inf_images, dates = dataset._load_images(path)
        
        assert len(vis_images) == 3
        assert len(inf_images) == 3
        assert dates is None


class TestForestNetNonGeoPlotting:
    """Test plotting functionality."""
    
    def test_basic_plot(self, temp_forestnet_dataset):
        """Test basic plotting."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        sample = dataset[0]
        
        # Image is already torch.Tensor with shape (C, T, H, W) after default_transform
        
        fig = dataset.plot(sample)
        
        assert fig is not None
        # Should have 4 subplots (3 timesteps + 1 legend)
        assert len(fig.axes) == 4
    
    def test_plot_with_suptitle(self, temp_forestnet_dataset):
        """Test plotting with custom suptitle."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        sample = dataset[0]
        
        # Image is already torch.Tensor
        
        fig = dataset.plot(sample, suptitle="Test Title")
        
        assert fig is not None
        assert fig._suptitle is not None
    
    def test_plot_missing_rgb_bands(self, temp_forestnet_dataset):
        """Test that plotting without RGB bands raises error."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            bands=("NIR", "SWIR_1")  # No RGB bands
        )
        sample = dataset[0]
        # Image is already torch.Tensor
        
        with pytest.raises(ValueError, match="RGB bands"):
            dataset.plot(sample)
    
    def test_plot_with_numpy_input(self, temp_forestnet_dataset):
        """Test plotting with numpy array input."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        sample = dataset[0]
        
        # Keep as numpy but add channel dimension first
        sample["image"] = np.transpose(sample["image"], (3, 0, 1, 2))
        
        fig = dataset.plot(sample)
        
        assert fig is not None


class TestForestNetNonGeoEdgeCases:
    """Test edge cases and error handling."""
    
    def test_length_method(self, temp_forestnet_dataset):
        """Test __len__ method."""
        dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split="train")
        
        assert len(dataset) == 10
        assert dataset.__len__() == 10
    
    def test_all_splits_available(self, temp_forestnet_dataset):
        """Test that all splits can be loaded."""
        for split in ["train", "val", "test"]:
            dataset = ForestNetNonGeo(data_root=temp_forestnet_dataset, split=split)
            assert len(dataset) > 0
            assert dataset.split == split
    
    def test_band_sets_property(self, temp_forestnet_dataset):
        """Test BAND_SETS property."""
        assert "all" in ForestNetNonGeo.BAND_SETS
        assert "rgb" in ForestNetNonGeo.BAND_SETS
        assert len(ForestNetNonGeo.BAND_SETS["all"]) == 6
        assert len(ForestNetNonGeo.BAND_SETS["rgb"]) == 3
    
    def test_all_band_names(self, temp_forestnet_dataset):
        """Test all_band_names property."""
        expected_bands = ("RED", "GREEN", "BLUE", "NIR", "SWIR_1", "SWIR_2")
        assert ForestNetNonGeo.all_band_names == expected_bands
    
    def test_rgb_bands(self, temp_forestnet_dataset):
        """Test rgb_bands property."""
        expected_rgb = ("RED", "GREEN", "BLUE")
        assert ForestNetNonGeo.rgb_bands == expected_rgb
    
    def test_splits_property(self, temp_forestnet_dataset):
        """Test splits property."""
        assert ForestNetNonGeo.splits == ("train", "test", "val")
    
    def test_default_label_map(self, temp_forestnet_dataset):
        """Test default_label_map property."""
        expected_map = {
            "Plantation": 0,
            "Smallholder agriculture": 1,
            "Grassland shrubland": 2,
            "Other": 3,
        }
        assert ForestNetNonGeo.default_label_map == expected_map
    
    def test_iterate_over_dataset(self, temp_forestnet_dataset):
        """Test iterating over the dataset."""
        # Use full dataset to avoid stratified sampling complexity
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train"
        )
        
        count = 0
        for sample in dataset:
            assert "image" in sample
            assert "label" in sample
            count += 1
            if count >= 3:  # Just test a few iterations
                break
        
        assert count == 3
    
    def test_infrared_band_selection(self, temp_forestnet_dataset):
        """Test selecting only infrared bands."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            bands=("NIR", "SWIR_1", "SWIR_2")
        )
        sample = dataset[0]
        
        # Should have 3 infrared bands - shape is (C, T, H, W)
        assert sample["image"].shape[0] == 3
    
    def test_single_band_selection(self, temp_forestnet_dataset):
        """Test selecting a single band."""
        dataset = ForestNetNonGeo(
            data_root=temp_forestnet_dataset,
            split="train",
            bands=("RED",)
        )
        sample = dataset[0]
        
        # Should have only 1 band - shape is (C, T, H, W)
        assert sample["image"].shape[0] == 1
