backend/mini_data.py

from __future__ import annotations

import shutil
from dataclasses import dataclass
from pathlib import Path

import pandas as pd


class UnsupportedDatasetTypeError(ValueError):
    """Raised when an uploaded dataset extension is outside Epic 1 scope."""


@dataclass(frozen=True)
class DatasetSummary:
    row_count: int
    column_count: int
    columns: list[str]
    file_size_bytes: int
    data_type: str
    canonical_data_path: Path
    mini_data_path: Path
    source_path: Path | None = None


class DatasetNormalizer:
    csv_extensions = {".csv"}
    excel_extensions = {".xls", ".xlsx"}

    def __init__(self, mini_data_sample_rows: int, chunk_size_rows: int) -> None:
        self.mini_data_sample_rows = mini_data_sample_rows
        self.chunk_size_rows = chunk_size_rows

    def normalize(self, source_file: Path, data_dir: Path) -> DatasetSummary:
        data_dir.mkdir(parents=True, exist_ok=True)
        file_extension = source_file.suffix.lower()
        canonical_data_path = data_dir / "data.csv"

        if file_extension in self.csv_extensions:
            shutil.copyfile(source_file, canonical_data_path)
            source_path: Path | None = None
            data_type = "csv"
        elif file_extension in self.excel_extensions:
            source_path = data_dir / f"source{file_extension}"
            shutil.copyfile(source_file, source_path)
            self._convert_excel_to_csv(source_file=source_path, csv_path=canonical_data_path)
            data_type = "excel"
        else:
            raise UnsupportedDatasetTypeError(
                f"Unsupported dataset extension: {file_extension}"
            )

        row_count, column_names, sample_frame = self._read_csv_summary(
            csv_path=canonical_data_path
        )
        mini_data_path = data_dir / "mini_data.csv"
        self._write_mini_data(sample_frame=sample_frame, mini_data_path=mini_data_path)

        return DatasetSummary(
            row_count=row_count,
            column_count=len(column_names),
            columns=column_names,
            file_size_bytes=canonical_data_path.stat().st_size,
            data_type=data_type,
            canonical_data_path=canonical_data_path,
            mini_data_path=mini_data_path,
            source_path=source_path,
        )

    @staticmethod
    def _convert_excel_to_csv(source_file: Path, csv_path: Path) -> None:
        excel_frame = pd.read_excel(source_file, sheet_name=0)
        excel_frame.to_csv(csv_path, index=False)

    def _read_csv_summary(self, csv_path: Path) -> tuple[int, list[str], pd.DataFrame]:
        row_count = 0
        column_names: list[str] = []
        sample_chunks: list[pd.DataFrame] = []
        sampled_rows = 0

        for data_chunk in pd.read_csv(csv_path, chunksize=self.chunk_size_rows):
            if not column_names:
                column_names = list(data_chunk.columns)
            row_count += len(data_chunk)

            if sampled_rows < self.mini_data_sample_rows:
                remaining_sample_rows = self.mini_data_sample_rows - sampled_rows
                sample_chunk = data_chunk.head(remaining_sample_rows)
                sample_chunks.append(sample_chunk)
                sampled_rows += len(sample_chunk)

        sample_frame = (
            pd.concat(sample_chunks, ignore_index=True)
            if sample_chunks
            else pd.DataFrame(columns=column_names)
        )
        return row_count, column_names, sample_frame

    @staticmethod
    def _write_mini_data(sample_frame: pd.DataFrame, mini_data_path: Path) -> None:
        mini_data = sample_frame.describe(include="all").transpose()
        mini_data.to_csv(mini_data_path)