# camera_trap_dataset.py
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np
from pathlib import Path
from typing import List, Tuple, Optional
from PIL import Image
import json


class CameraTrapDataset(Dataset):
    """Dataset for camera trap image classification."""

    SPECIES_MAP = {
        "empty": 0,
        "deer": 1,
        "fox": 2,
        "rabbit": 3,
        "bird": 4,
        "wolf": 5,
        "bear": 6,
        "elk": 7,
        "bobcat": 8,
        "cougar": 9,
        "lynx": 10,
        "eagle": 11,
        "owl": 12,
        "squirrel": 13,
        "raccoon": 14,
        "skunk": 15,
        "porcupine": 16,
        "beaver": 17,
        "muskrat": 18,
        "human": 19,
        "vehicle": 20,
    }

    def __init__(
        self,
        image_dir: str,
        annotations_file: str,
        split: str = "train",
        image_size: int = 224,
    ):
        self.image_dir = Path(image_dir)
        self.image_size = image_size
        self.split = split

        with open(annotations_file) as f:
            self.annotations = json.load(f)

        self.image_paths = []
        self.labels = []
        for ann in self.annotations:
            img_path = self.image_dir / ann["filename"]
            if img_path.exists():
                self.image_paths.append(img_path)
                self.labels.append(self.SPECIES_MAP.get(ann["species"], 0))

        self._compute_class_weights()

        if split == "train":
            self.transform = A.Compose([
                A.Resize(image_size, image_size),
                A.HorizontalFlip(p=0.5),
                A.RandomBrightnessContrast(p=0.3),
                A.RandomRotate90(p=0.2),
                A.GaussNoise(var_limit=(10, 50), p=0.2),
                A.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                ),
                ToTensorV2(),
            ])
        else:
            self.transform = A.Compose([
                A.Resize(image_size, image_size),
                A.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                ),
                ToTensorV2(),
            ])

    def _compute_class_weights(self):
        """Compute inverse frequency class weights."""
        class_counts = np.bincount(self.labels, minlength=len(self.SPECIES_MAP))
        class_counts = np.maximum(class_counts, 1)
        total = sum(class_counts)
        self.class_weights = total / (len(class_counts) * class_counts)
        self.class_weights = torch.tensor(self.class_weights, dtype=torch.float)

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        image = np.array(Image.open(self.image_paths[idx]).convert("RGB"))
        label = self.labels[idx]

        transformed = self.transform(image=image)
        image_tensor = transformed["image"]

        return image_tensor, label

    def get_weighted_sampler(self) -> torch.utils.data.WeightedRandomSampler:
        """Create weighted sampler for class imbalance."""
        sample_weights = [self.class_weights[label] for label in self.labels]
        return torch.utils.data.WeightedRandomSampler(
            weights=sample_weights,
            num_samples=len(sample_weights),
            replacement=True,
        )

# detection_model.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
from typing import Tuple, Optional
import timm


class AnimalDetector(nn.Module):
    """YOLO-based animal detector for camera traps."""

    def __init__(self, num_classes: int = 21, pretrained: bool = True):
        super().__init__()
        from ultralytics import YOLO

        self.yolo = YOLO("yolov8n.pt")
        self.num_classes = num_classes

    def forward(self, x: torch.Tensor) -> dict:
        results = self.yolo(x)
        return results


class SpeciesClassifier(nn.Module):
    """EfficientNet-based species classifier."""

    def __init__(
        self,
        num_classes: int = 21,
        model_name: str = "efficientnet_b3",
        pretrained: bool = True,
        dropout: float = 0.3,
    ):
        super().__init__()
        self.backbone = timm.create_model(
            model_name, pretrained=pretrained, num_classes=0
        )
        feature_dim = self.backbone.num_features

        self.classifier = nn.Sequential(
            nn.Linear(feature_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(256, num_classes),
        )

        self.attention = nn.Sequential(
            nn.Linear(feature_dim, 128),
            nn.Tanh(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        features = self.backbone(x)

        attn_weights = self.attention(features)
        features = features * attn_weights

        logits = self.classifier(features)
        return logits

    def predict_with_confidence(
        self, x: torch.Tensor, threshold: float = 0.5
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Predict with confidence filtering."""
        logits = self.forward(x)
        probs = F.softmax(logits, dim=1)
        confidence, predictions = probs.max(dim=1)

        predictions[confidence < threshold] = 0

        return predictions, confidence, probs


class CameraTrapPipeline(nn.Module):
    """Full detection + classification pipeline."""

    def __init__(self, num_classes: int = 21):
        super().__init__()
        self.detector = AnimalDetector()
        self.classifier = SpeciesClassifier(num_classes=num_classes)

    def forward(self, images: torch.Tensor) -> dict:
        detections = self.detector(images)

        crops = []
        for det in detections:
            boxes = det.boxes
            if boxes is not None and len(boxes) > 0:
                for box in boxes.xyxy:
                    x1, y1, x2, y2 = box.int()
                    crop = images[:, :, y1:y2, x1:x2]
                    crop = F.interpolate(
                        crop, size=(224, 224), mode="bilinear"
                    )
                    crops.append(crop)

        if crops:
            crops = torch.cat(crops, dim=0)
            species_logits = self.classifier(crops)
            species_probs = F.softmax(species_logits, dim=1)
        else:
            species_probs = torch.zeros(len(images), self.classifier.classifier[-1].out_features)

        return {
            "detections": detections,
            "species_logits": species_logits if crops is not None else None,
            "species_probs": species_probs,
        }

# population_analysis.py
import numpy as np
from typing import List, Dict, Tuple
from dataclasses import dataclass
from datetime import datetime, timedelta
import json


@dataclass
class Observation:
    """Single species observation."""
    timestamp: datetime
    species: str
    individual_id: Optional[str]
    location: Tuple[float, float]
    confidence: float
    camera_id: str


class PopulationAnalyzer:
    """Analyze species populations from camera trap observations."""

    def __init__(self, observations: List[Observation]):
        self.observations = observations

    def species_abundance(self, time_period: str = "monthly") -> Dict:
        """Calculate species abundance indices."""
        species_counts = {}
        for obs in self.observations:
            species_counts.setdefault(obs.species, []).append(obs)

        abundance = {}
        for species, obs_list in species_counts.items():
            individuals = set(obs.individual_id for obs in obs_list if obs.individual_id)
            abundance[species] = {
                "total_observations": len(obs_list),
                "unique_individuals": len(individuals),
                "capture_rate": len(obs_list) / max(1, len(set(o.camera_id for o in obs_list))),
            }

        return abundance

    def occupancy_model(
        self, species: str, time_windows: int = 12
    ) -> Dict:
        """Estimate occupancy probability."""
        species_obs = [o for o in self.observations if o.species == species]
        cameras = set(o.camera_id for o in self.observations)

        occupancy_data = []
        window_duration = timedelta(days=30)

        for i in range(time_windows):
            start = datetime(2024, 1, 1) + i * window_duration
            end = start + window_duration

            occupied = set()
            for obs in species_obs:
                if start <= obs.timestamp < end:
                    occupied.add(obs.camera_id)

            occupancy_rate = len(occupied) / len(cameras) if cameras else 0
            occupancy_data.append({
                "period": i,
                "start": start.isoformat(),
                "occupancy_rate": occupancy_rate,
                "cameras_occupied": len(occupied),
                "total_cameras": len(cameras),
            })

        return {
            "species": species,
            "mean_occupancy": np.mean([d["occupancy_rate"] for d in occupancy_data]),
            "occupancy_trend": self._compute_trend(
                [d["occupancy_rate"] for d in occupancy_data]
            ),
            "period_data": occupancy_data,
        }

    def activity_patterns(
        self, species: str, time_resolution: int = 1
    ) -> Dict:
        """Analyze diel activity patterns."""
        species_obs = [o for o in self.observations if o.species == species]

        hourly_counts = np.zeros(24)
        for obs in species_obs:
            hour = obs.timestamp.hour
            hourly_counts[hour] += 1

        hourly_probs = hourly_counts / max(hourly_counts.sum(), 1)

        nocturnal_index = hourly_probs[20:24].sum() + hourly_probs[0:6].sum()

        return {
            "species": species,
            "hourly_distribution": hourly_probs.tolist(),
            "nocturnal_index": nocturnal_index,
            "peak_hour": int(np.argmax(hourly_probs)),
            "activity_type": "nocturnal" if nocturnal_index > 0.6 else "diurnal",
        }

    def population_trend(
        self, species: str, method: str = "capture_recapture"
    ) -> Dict:
        """Estimate population trend using capture-recapture."""
        species_obs = sorted(
            [o for o in self.observations if o.species == species],
            key=lambda x: x.timestamp,
        )

        if not species_obs:
            return {"species": species, "trend": "insufficient_data"}

        periods = []
        period_size = timedelta(days=90)

        start_date = species_obs[0].timestamp
        end_date = species_obs[-1].timestamp

        current = start_date
        while current < end_date:
            period_obs = [
                o for o in species_obs
                if current <= o.timestamp < current + period_size
            ]

            if period_obs:
                individuals = set(o.individual_id for o in period_obs if o.individual_id)
                periods.append({
                    "start": current.isoformat(),
                    "count": len(individuals),
                    "observations": len(period_obs),
                })

            current += period_size

        if len(periods) < 2:
            return {"species": species, "trend": "insufficient_data"}

        counts = [p["count"] for p in periods]
        trend = self._compute_trend(counts)

        return {
            "species": species,
            "method": method,
            "periods": periods,
            "mean_population": np.mean(counts),
            "trend_direction": "increasing" if trend > 0.1 else "decreasing" if trend < -0.1 else "stable",
            "trend_strength": abs(trend),
        }

    def _compute_trend(self, values: List[float]) -> float:
        """Compute normalized linear trend."""
        if len(values) < 2:
            return 0.0
        x = np.arange(len(values))
        coeffs = np.polyfit(x, values, 1)
        return coeffs[0] / max(np.mean(values), 1e-6)