main.py

# !pip install torch torchvision numpy pandas loguru astropy astroquery tensorboard

"""
SETI AI Framework
A production-grade implementation for AI-driven analysis of astronomical data for SETI research.
"""

from datetime import datetime, timedelta
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from astropy.io import fits
from astroquery.mast import Observations
from loguru import logger
from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader, Dataset
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms

# Configure logging
logger.add(
    "logs/seti_framework_{time}.log",
    rotation="100 MB",
    retention="30 days",
    level="INFO",
)


class Config:
    """Configuration settings for the SETI AI Framework."""

    # Data settings
    BASE_PATH = Path("seti_data")
    MODEL_PATH = BASE_PATH / "models"
    DATA_PATH = BASE_PATH / "raw_data"
    PROCESSED_PATH = BASE_PATH / "processed"

    # Training settings
    BATCH_SIZE = 64
    LEARNING_RATE = 1e-4
    NUM_EPOCHS = 100
    DEVICE = torch.device(
        "cuda" if torch.cuda.is_available() else "cpu"
    )

    # Model architecture
    INPUT_SIZE = 1024  # Size of input spectrogram
    HIDDEN_SIZE = 512
    NUM_CLASSES = 2  # Binary classification (signal vs. no signal)

    def __init__(self):
        # Create necessary directories
        for path in [
            self.BASE_PATH,
            self.MODEL_PATH,
            self.DATA_PATH,
            self.PROCESSED_PATH,
        ]:
            path.mkdir(parents=True, exist_ok=True)


class DataFetcher:
    """Handles fetching and preprocessing of astronomical data."""

    def __init__(self, config: Config):
        self.config = config
        logger.info("Initializing DataFetcher")

    def fetch_radio_data(
        self, start_date: datetime, end_date: datetime
    ) -> List[Path]:
        """
        Fetches radio telescope data from various sources.

        Args:
            start_date: Start date for data collection
            end_date: End date for data collection

        Returns:
            List of paths to downloaded data files
        """
        logger.info(
            f"Fetching radio data from {start_date} to {end_date}"
        )

        try:
            # Query Breakthrough Listen Open Data Archive
            # Note: Replace with actual API endpoint
            observations = Observations.query_criteria(
                dataproduct_type="spectrum",
                t_min=start_date.strftime("%Y-%m-%d"),
                t_max=end_date.strftime("%Y-%m-%d"),
            )

            downloaded_files = []
            for obs in observations:
                file_path = (
                    self.config.DATA_PATH / f"{obs['obs_id']}.fits"
                )
                if not file_path.exists():
                    products = Observations.get_product_list(obs)
                    Observations.download_products(
                        products,
                        download_dir=str(self.config.DATA_PATH),
                    )
                downloaded_files.append(file_path)

            return downloaded_files

        except Exception as e:
            logger.error(f"Error fetching radio data: {str(e)}")
            raise


class SignalProcessor:
    """Processes raw astronomical data into ML-ready format."""

    def __init__(self, config: Config):
        self.config = config
        self.scaler = StandardScaler()
        logger.info("Initializing SignalProcessor")

    def process_file(self, file_path: Path) -> np.ndarray:
        """
        Processes a single FITS file into a spectrogram.

        Args:
            file_path: Path to FITS file

        Returns:
            Processed spectrogram as numpy array
        """
        try:
            with fits.open(file_path) as hdul:
                data = hdul[0].data

            # Basic preprocessing steps
            data = self._remove_background(data)
            data = self._normalize(data)
            spectrogram = self._create_spectrogram(data)

            return spectrogram

        except Exception as e:
            logger.error(
                f"Error processing file {file_path}: {str(e)}"
            )
            raise

    def _remove_background(self, data: np.ndarray) -> np.ndarray:
        """Removes background noise from the data."""
        background = np.median(data, axis=0)
        return data - background

    def _normalize(self, data: np.ndarray) -> np.ndarray:
        """Normalizes the data using StandardScaler."""
        shape = data.shape
        flat_data = data.reshape(-1, 1)
        normalized = self.scaler.fit_transform(flat_data)
        return normalized.reshape(shape)

    def _create_spectrogram(self, data: np.ndarray) -> np.ndarray:
        """Creates a spectrogram from the time series data."""
        return np.abs(np.fft.rfft2(data))


class SETIDataset(Dataset):
    """PyTorch dataset for SETI data."""

    def __init__(
        self,
        spectrograms: List[np.ndarray],
        labels: List[int],
        transform: Optional[transforms.Compose] = None,
    ):
        self.spectrograms = spectrograms
        self.labels = labels
        self.transform = transform

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

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        spectrogram = self.spectrograms[idx]
        label = self.labels[idx]

        if self.transform:
            spectrogram = self.transform(spectrogram)

        return torch.FloatTensor(spectrogram), label


class SETINet(nn.Module):
    """Neural network architecture for SETI signal detection."""

    def __init__(self, config: Config):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

        self.classifier = nn.Sequential(
            nn.Linear(
                128
                * (config.INPUT_SIZE // 8)
                * (config.INPUT_SIZE // 8),
                config.HIDDEN_SIZE,
            ),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(config.HIDDEN_SIZE, config.NUM_CLASSES),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x


class Trainer:
    """Handles model training and evaluation."""

    def __init__(self, config: Config):
        self.config = config
        self.device = config.DEVICE
        self.writer = SummaryWriter(log_dir="runs/seti_experiment")
        logger.info(
            f"Initializing Trainer with device: {self.device}"
        )

    def train(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
    ) -> nn.Module:
        """
        Trains the model on the provided data.

        Args:
            model: Neural network model
            train_loader: Training data loader
            val_loader: Validation data loader

        Returns:
            Trained model
        """
        model = model.to(self.device)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(
            model.parameters(), lr=self.config.LEARNING_RATE
        )

        best_val_loss = float("inf")

        for epoch in range(self.config.NUM_EPOCHS):
            logger.info(f"Starting epoch {epoch + 1}")

            # Training phase
            model.train()
            train_loss = 0.0
            for batch_idx, (data, target) in enumerate(train_loader):
                data, target = data.to(self.device), target.to(
                    self.device
                )

                optimizer.zero_grad()
                output = model(data)
                loss = criterion(output, target)
                loss.backward()
                optimizer.step()

                train_loss += loss.item()

                if batch_idx % 10 == 0:
                    logger.info(
                        f"Batch {batch_idx}: Loss {loss.item():.4f}"
                    )

            # Validation phase
            model.eval()
            val_loss = 0.0
            correct = 0
            total = 0

            with torch.no_grad():
                for data, target in val_loader:
                    data, target = data.to(self.device), target.to(
                        self.device
                    )
                    output = model(data)
                    val_loss += criterion(output, target).item()

                    _, predicted = torch.max(output.data, 1)
                    total += target.size(0)
                    correct += (predicted == target).sum().item()

            val_accuracy = 100 * correct / total
            val_loss /= len(val_loader)

            # Log metrics
            self.writer.add_scalar(
                "Loss/train", train_loss / len(train_loader), epoch
            )
            self.writer.add_scalar("Loss/validation", val_loss, epoch)
            self.writer.add_scalar(
                "Accuracy/validation", val_accuracy, epoch
            )

            logger.info(
                f"Epoch {epoch + 1}: "
                f"Train Loss {train_loss / len(train_loader):.4f}, "
                f"Val Loss {val_loss:.4f}, "
                f"Val Accuracy {val_accuracy:.2f}%"
            )

            # Save best model
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                self.save_model(model, epoch, val_loss)

        self.writer.close()
        return model

    def save_model(
        self, model: nn.Module, epoch: int, val_loss: float
    ) -> None:
        """Saves model checkpoint."""
        checkpoint_path = (
            self.config.MODEL_PATH
            / f"model_epoch_{epoch}_loss_{val_loss:.4f}.pt"
        )
        torch.save(
            {
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "val_loss": val_loss,
            },
            checkpoint_path,
        )
        logger.info(f"Saved model checkpoint to {checkpoint_path}")


def main():
    """Main execution function."""
    try:
        # Initialize configuration
        config = Config()
        logger.info("Starting SETI AI Framework")

        # Initialize components
        data_fetcher = DataFetcher(config)
        signal_processor = SignalProcessor(config)

        # Fetch data for the past 10 years
        end_date = datetime.now()
        start_date = end_date - timedelta(days=365 * 10)

        data_files = data_fetcher.fetch_radio_data(
            start_date, end_date
        )
        logger.info(f"Retrieved {len(data_files)} data files")

        # Process data
        spectrograms = []
        labels = []  # In practice, you'd need a labeling strategy

        for file_path in data_files:
            spectrogram = signal_processor.process_file(file_path)
            spectrograms.append(spectrogram)
            # Placeholder for labeling strategy
            labels.append(np.random.randint(0, 2))

        # Create datasets
        dataset = SETIDataset(
            spectrograms=spectrograms,
            labels=labels,
            transform=transforms.Compose(
                [
                    transforms.ToTensor(),
                    transforms.Normalize(mean=[0.5], std=[0.5]),
                ]
            ),
        )

        # Split dataset
        train_size = int(0.8 * len(dataset))
        val_size = len(dataset) - train_size
        train_dataset, val_dataset = torch.utils.data.random_split(
            dataset, [train_size, val_size]
        )

        # Create data loaders
        train_loader = DataLoader(
            train_dataset,
            batch_size=config.BATCH_SIZE,
            shuffle=True,
            num_workers=4,
        )
        val_loader = DataLoader(
            val_dataset,
            batch_size=config.BATCH_SIZE,
            shuffle=False,
            num_workers=4,
        )

        # Initialize and train model
        model = SETINet(config)
        trainer = Trainer(config)

        trainer.train(model, train_loader, val_loader)
        logger.info("Training completed successfully")

    except Exception as e:
        logger.error(f"Error in main execution: {str(e)}")
        raise


if __name__ == "__main__":
    main()