from eotorch.data import get_segmentation_dataset

# Create a dataset with images and labels
ds = get_segmentation_dataset(
    images_dir="path/to/images",
    labels_dir="path/to/labels",
    all_image_bands=("B", "G", "R", "NIR"),  # Band names in your files
    rgb_bands=("R", "G", "B"),                # Bands for visualization
    class_mapping={1: "Water", 2: "Forest", 3: "Urban"},
    reduce_zero_label=True,  # Shift labels (1,2,3) → (0,1,2) for training
)

# Visualize samples to verify correct image-label pairing
ds.plot_samples(n=3, patch_size=256, show_filepaths=True)

ds = get_segmentation_dataset(
    images_dir="path/to/images",
    labels_dir="path/to/labels",
    # Match files by date extracted from filename
    image_filename_regex=r"(?P<date>\d{8})_S2[AB]_.*\.tif",
    label_filename_regex=r"(?P<date>\d{8})_.*\.tif",
    image_date_format="%Y%m%d",
    label_date_format="%Y%m%d",
    # ...
)

image_filename_regex=r"(?P<date>\d{8})_S2[AB]_\w+_\d{8}_\d+_L2A.*\.tif"
label_filename_regex=r"(?P<date>\d{8})_labels\.tif"
image_date_format="%Y%m%d"
label_date_format="%Y%m%d"

image_filename_regex=r"satellite_(?P<date>\d{4}-\d{2}-\d{2})_.*\.tif"
label_filename_regex=r"mask_(?P<date>\d{4}-\d{2}-\d{2})\.tif"
image_date_format="%Y-%m-%d"
label_date_format="%Y-%m-%d"

image_filename_regex=r"landcover_(?P<date>\d{4})_.*\.tif"
label_filename_regex=r"training_(?P<date>\d{4})_.*\.tif"
image_date_format="%Y"
label_date_format="%Y"

from eotorch.data import splits

# Split 80% train, 20% validation (randomly assigned)
train_ds, val_ds = splits.file_wise_split(
    dataset=ds, 
    ratios_or_counts=[0.8, 0.2]
)

# Three-way split: 70% train, 15% val, 15% test
train_ds, val_ds, test_ds = splits.file_wise_split(
    dataset=ds,
    ratios_or_counts=[0.7, 0.15, 0.15]
)

# Explicitly assign files to validation set (remaining go to train)
train_ds, val_ds = splits.file_wise_split(
    dataset=ds,
    val_img_files=["scene1.tif", "scene2.tif"]
)

# Specify both validation and test files
train_ds, val_ds, test_ds = splits.file_wise_split(
    dataset=ds,
    val_img_files=["val_scene1.tif", "val_scene2.tif"],
    test_img_files=["test_scene1.tif"]
)

# Or use glob patterns to select files
train_ds, val_ds = splits.file_wise_split(
    dataset=ds,
    val_img_glob="/path/to/data/*_2024_*.tif"  # All 2024 scenes for validation
)

train_ds, val_ds = splits.aoi_split(
    dataset=ds,
    aoi_files="validation_region.geojson",
    buffer_size_metres=100,  # Optional buffer to prevent spatial leakage
)

from eotorch.data import SegmentationDataModule

module = SegmentationDataModule(
    train_dataset=train_ds,
    val_dataset=val_ds,
    batch_size=16,
    patch_size=224,         # Size of patches sampled from images
    num_workers=4,
    pin_memory=True,
)

module = SegmentationDataModule(
    train_dataset=train_ds,
    # Default: GridGeoSampler with stride=patch_size/2
    # For random sampling:
    train_sampler_config={"type": "RandomGeoSampler", "length": 5000},
    # ...
)

from eotorch.data import SemanticSegmentationTask

result_path = SemanticSegmentationTask.predict_on_tif_file(
    tif_file_path="large_image.tif",
    checkpoint_path="checkpoints/best-model.ckpt",
    batch_size=8,
    show_results=True,        # Display preview in notebook
    out_file_path="predictions/output.tif",  # Optional, auto-generated if not set
)

from eotorch.processing import zscore_normalize

# Computes mean/std across all images and applies: (x - mean) / std
stats = zscore_normalize(
    input_files="path/to/images",      # Directory, file, or list of files
    out_dir="path/to/normalized",
    save_stats=True,                   # Saves mean/std to JSON for inference
    sample_size=1.0,                   # Use < 1.0 to sample for speed on large datasets
)
# Output files have "_zscore" suffix

from eotorch.processing import normalize

# Clips to percentile bounds and scales to [0, 1]
normalize(
    img_path="image.tif",
    limits=(1, 99),            # Percentile cutoffs (or dict of per-band limits)
    out_path="image_norm.tif",
    out_res=10.0,              # Optional: resample to target resolution
)

from eotorch.processing.rasterize import FileSource

# Create rasterizer with class files aligned to a reference image
source = FileSource(
    classes=["water", "forest", "urban"],  # Looks for water.shp, forest.shp, etc.
    root_dir="path/to/shapefiles",
    img_path="reference_image.tif"         # Matches CRS, resolution, and extent
)

# Generate raster labels (class 1=water, 2=forest, 3=urban)
source.rasterize_polygons(
    out_path="labels.tif",
    exclude_nodata_bounds=True,    # Crop to data extent
    class_order=[3, 2, 1],         # Optional: control overlap priority
)

from rasterio.crs import CRS
from affine import Affine

source = FileSource(
    classes=["water", "forest", "urban"],
    root_dir="path/to/shapefiles",
    transform=Affine(10.0, 0.0, 500000.0, 0.0, -10.0, 6000000.0),  # 10m resolution
    shape=(1024, 1024),                                            # Height x Width
    crs=CRS.from_epsg(32632)                                       # UTM zone 32N
)

from eotorch.processing.rasterize import DataframeSource
import geopandas as gpd

gdf = gpd.read_file("annotations.geojson")

source = DataframeSource(
    dataframe=gdf,
    class_column="class_id",       # Column containing integer class labels
    img_path="reference_image.tif"
)

source.rasterize_polygons(out_path="labels.tif")

# TorchGeo requires passing the class, not an instance
datamodule = SomeGeoDataModule(
    dataset_class=MyRasterDataset,  # Class type
    batch_size=8,
    **dataset_kwargs  # Args passed to class constructor
)

# EOTorch lets you create and inspect first
ds = get_segmentation_dataset(...)
ds.plot_samples(n=3)  # Verify everything looks correct

train_ds, val_ds = splits.file_wise_split(ds, [0.8, 0.2])

# Then pass instances to DataModule
module = SegmentationDataModule(
    train_dataset=train_ds,  # Instance
    val_dataset=val_ds,      # Instance
)

from terratorch.datasets import HLSBands
from terratorch.tasks import SemanticSegmentationTask
from lightning.pytorch import Trainer

# Define bands (must match your data)
BANDS = [HLSBands.RED, HLSBands.GREEN, HLSBands.BLUE]

# Create the model task
model = SemanticSegmentationTask(
    model_args=dict(
        backbone="prithvi_eo_v2_300",
        backbone_pretrained=True,
        backbone_bands=BANDS,
        backbone_num_frames=1,
        decoder="UperNetDecoder",
        decoder_channels=256,
        num_classes=3,
        necks=[
            {"name": "SelectIndices", "indices": [5, 11, 17, 23]},
            {"name": "ReshapeTokensToImage"},
        ],
    ),
    loss="focal",
    lr=1e-4,
    optimizer="AdamW",
    ignore_index=-1,
    freeze_backbone=True,
    model_factory="EncoderDecoderFactory",
)

# Train with EOTorch DataModule
trainer = Trainer(max_epochs=50, accelerator="auto")
trainer.fit(model, datamodule=module)

import torch
from terratorch.tasks import SemanticSegmentationTask
from eotorch.inference.inference import predict_on_tif_generic

# Load the trained checkpoint
checkpoint_path = "results/Prithvi/baseline/checkpoints/best-checkpoint.ckpt"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = SemanticSegmentationTask.load_from_checkpoint(checkpoint_path, map_location=device)
model.to(device)
model.eval()

# Define a prediction function compatible with predict_on_tif_generic
def predict_func(batch):
    """Prediction function for TerraTorch models."""
    with torch.no_grad():
        if not isinstance(batch, torch.Tensor):
            batch = torch.tensor(batch, dtype=torch.float32)
        batch = batch.to(device)
        
        # TerraTorch models return a ModelOutput object
        logits = model.model(batch)
        logits = logits.output  # Extract the actual tensor
        
        predictions = torch.argmax(logits, dim=1)
        
        # If reduce_zero_label was used during training, shift predictions back
        predictions += 1  # Adjust if your labels were shifted
        
    return predictions

# Run inference on a large raster
result = predict_on_tif_generic(
    tif_file_path="path/to/input_image.tif",
    prediction_func=predict_func,
    patch_size=224,              # Must match training patch size
    batch_size=8,
    out_file_path="predictions/output.tif",
    show_results=True,           # Display preview in notebook
    class_mapping=class_mapping, # For visualization
)

from pathlib import Path
from lightning.pytorch import Trainer
from eotorch.data import get_segmentation_dataset, SegmentationDataModule, splits

# 1. Create dataset
ds = get_segmentation_dataset(
    images_dir=Path("data/images"),
    labels_dir=Path("data/labels"),
    all_image_bands=("R", "G", "B"),
    rgb_bands=("R", "G", "B"),
    class_mapping={1: "Class A", 2: "Class B"},
    reduce_zero_label=True,
)

# 2. Verify with visualization
ds.plot_samples(n=2, patch_size=256)

# 3. Split data
train_ds, val_ds = splits.file_wise_split(ds, ratios_or_counts=[0.8, 0.2])

# 4. Create data module
module = SegmentationDataModule(
    train_dataset=train_ds,
    val_dataset=val_ds,
    batch_size=16,
    patch_size=224,
    num_workers=4,
)

# 5. Train (using your Lightning module)
trainer = Trainer(max_epochs=10)
trainer.fit(your_model, datamodule=module)