"""
Experiment results collection and export for plate reader data
"""
from .cdcl import CDCLParser, PlateReaderModeCalculator
from .cdcl.codes import GenericResponsePrefixes
from .utils import index_to_well
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any, Callable
from datetime import datetime
import json
import csv
import io
from enum import Enum


class DataFormat(Enum):
    """Export data formats"""
    CSV = "csv"
    JSON = "json"
    DICT = "dict"


@dataclass
class DataPoint:
    """Single measurement data point"""
    timestamp: int  # Unix timestamp
    well: str  # Well identifier (e.g., "A1", "B2")
    wavelength: int  # Wavelength in nm
    value: float  # OD or fluorescence reading
    temperature: Optional[float] = None  # Temperature at time of reading
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'timestamp': self.timestamp,
            'datetime': datetime.fromtimestamp(self.timestamp).isoformat(),
            'well': self.well,
            'wavelength': self.wavelength,
            'value': self.value,
            'temperature': self.temperature,
            **self.metadata
        }


class ExperimentResults:
    """
    Container for experiment results with export capabilities.

    Collects data points during an experiment and provides methods
    for exporting to various formats.
    """

    def __init__(
        self,
        experiment_id: str,
        experiment_name: str = "",
        wavelengths: Optional[List[int]] = None,
        plate_layout: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize experiment results container.

        Args:
            experiment_id: Unique experiment identifier
            experiment_name: Human-readable experiment name
            wavelengths: List of wavelengths being measured
            plate_layout: Dictionary describing plate layout (rows, cols, etc.)
        """
        self.experiment_id = experiment_id
        self.experiment_name = experiment_name
        self.wavelengths = wavelengths or []
        self.plate_layout = plate_layout or {}

        # Store data points (individual well OD reading)
        self._data_points: List[DataPoint] = []

        # Metadata
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self.metadata: Dict[str, Any] = {}

        # Callbacks for real-time processing
        self._data_callbacks: List[Callable[[DataPoint], None]] = []

    def add_data_point(
        self,
        timestamp: int,
        well: str,
        wavelength: int,
        value: float,
        temperature: Optional[float] = None,
        **metadata
    ) -> None:
        """
        Add a single data point.

        Args:
            timestamp: Unix timestamp of measurement
            well: Well identifier
            wavelength: Wavelength in nm
            value: OD or fluorescence value
            temperature: Optional temperature reading
            **metadata: Additional metadata
        """
        data_point = DataPoint(
            timestamp=timestamp,
            well=well,
            wavelength=wavelength,
            value=value,
            temperature=temperature,
            metadata=metadata
        )

        self._data_points.append(data_point)

        # Update start/end times
        if self.start_time is None or timestamp < self.start_time:
            self.start_time = timestamp
        if self.end_time is None or timestamp > self.end_time:
            self.end_time = timestamp

        # Trigger callbacks
        for callback in self._data_callbacks:
            callback(data_point)

    def add_time_point(
        self,
        timestamp: int,
        measurements: Dict[str, float],
        wavelength: int,
        temperature: Optional[float] = None,
        **metadata
    ) -> None:
        """
        Add all measurements from a single time point.

        Args:
            timestamp: Unix timestamp
            measurements: Dictionary of well -> value
            wavelength: Wavelength in nm
            temperature: Optional temperature
            **metadata: Additional metadata
        """

        # Also add individual data points
        for well, value in measurements.items():
            self.add_data_point(
                timestamp=timestamp,
                well=well,
                wavelength=wavelength,
                value=value,
                temperature=temperature,
                **metadata
            )

    def register_callback(self, callback: Callable[[DataPoint], None]) -> None:
        """
        Register a callback for real-time data processing.

        Args:
            callback: Function that takes a DataPoint
        """
        self._data_callbacks.append(callback)

    def get_data_points(
        self,
        well: Optional[str] = None,
        wavelength: Optional[int] = None
    ) -> List[DataPoint]:
        """
        Get filtered data points.

        Args:
            well: Filter by well (None for all)
            wavelength: Filter by wavelength (None for all)

        Returns:
            List of matching data points
        """
        points = self._data_points

        if well is not None:
            points = [p for p in points if p.well == well]

        if wavelength is not None:
            points = [p for p in points if p.wavelength == wavelength]

        return points

    def get_well_timeseries(
        self,
        well: str,
        wavelength: Optional[int] = None
    ) -> List[tuple[float, float]]:
        """
        Get time series data for a specific well.

        Args:
            well: Well identifier
            wavelength: Optional wavelength filter

        Returns:
            List of (timestamp, value) tuples sorted by time
        """
        points = self.get_data_points(well=well, wavelength=wavelength)
        timeseries = [(p.timestamp, p.value) for p in points]
        timeseries.sort(key=lambda x: x[0])
        return timeseries

    def __sort_wells(self, wells):

        def sort_key(well):
            # Split into row (letters) and column (numbers)
            row = well[0].upper()        # e.g. 'A'
            col = int(well[1:])          # e.g. 1, 2, 12
            return (row, col)

        return sorted(wells, key=sort_key)

    def get_wells(self) -> List[str]:
        """Get list of all wells with data"""
        wells = set(p.well for p in self._data_points)
        return self.__sort_wells(wells)

    def get_wavelengths(self) -> List[int]:
        """Get list of all wavelengths measured"""
        wavelengths = set(p.wavelength for p in self._data_points)
        return sorted(wavelengths)

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert results to dictionary.

        Returns:
            Dictionary containing all experiment data
        """
        return {
            'experiment_id': self.experiment_id,
            'experiment_name': self.experiment_name,
            'start_time': self.start_time,
            'end_time': self.end_time,
            'start_datetime': datetime.fromtimestamp(self.start_time).isoformat() if self.start_time else None,
            'end_datetime': datetime.fromtimestamp(self.end_time).isoformat() if self.end_time else None,
            'duration_seconds': (self.end_time - self.start_time) if self.start_time and self.end_time else None,
            'wavelengths': self.wavelengths,
            'plate_layout': self.plate_layout,
            'metadata': self.metadata,
            'data_points': [p.to_dict() for p in self._data_points],
            'total_measurements': len(self._data_points),
            'wells_measured': self.get_wells(),
        }

    def to_json(self, indent: int = 2) -> str:
        """
        Export results as JSON string.

        Args:
            indent: JSON indentation level

        Returns:
            JSON string
        """
        return json.dumps(self.to_dict(), indent=indent)

    def to_csv(
        self,
        format: str = "long",
        include_metadata: bool = True
    ) -> str:
        """
        Export results as CSV string.

        Args:
            format: "long" for one row per measurement,
                   "wide" for one row per timepoint with well columns
            include_metadata: Include temperature and other metadata

        Returns:
            CSV string
        """
        output = io.StringIO()

        if format == "long":
            # Long format: one row per measurement
            writer = csv.writer(output)

            # Header
            header = ['timestamp', 'datetime', 'well', 'wavelength', 'value']
            if include_metadata:
                header.extend(['temperature'])
            writer.writerow(header)

            # Data rows
            for point in self._data_points:
                row = [
                    point.timestamp,
                    datetime.fromtimestamp(point.timestamp).isoformat(),
                    point.well,
                    point.wavelength,
                    point.value
                ]
                if include_metadata:
                    row.append(point.temperature)
                writer.writerow(row)

        elif format == "wide":
            # Wide format: one row per timepoint, columns for each well
            writer = csv.writer(output)

            # Get all wells and wavelengths
            wells = self.get_wells()
            wavelengths = self.get_wavelengths()

            # Group by wavelength
            for wavelength in wavelengths:
                # Header
                header = ['timestamp', 'datetime']
                if include_metadata:
                    header.append('temperature')
                header.extend(wells)
                writer.writerow(header)

                # Group data by timestamp
                time_data: Dict[float, Dict[str, Any]] = {}
                for point in self._data_points:
                    if point.wavelength == wavelength:
                        if point.timestamp not in time_data:
                            time_data[point.timestamp] = {
                                'temperature': point.temperature,
                                'wells': {}
                            }
                        time_data[point.timestamp]['wells'][point.well] = point.value

                # Write rows
                for timestamp in sorted(time_data.keys()):
                    data = time_data[timestamp]
                    row = [
                        timestamp,
                        datetime.fromtimestamp(timestamp).isoformat()
                    ]
                    if include_metadata:
                        row.append(data['temperature'])

                    # Add well values (empty string if missing)
                    for well in wells:
                        row.append(data['wells'].get(well, ''))

                    writer.writerow(row)

                # Blank line between wavelengths
                if len(wavelengths) > 1:
                    writer.writerow([])

        else:
            raise ValueError(f"Unknown format: {format}. Use 'long' or 'wide'")

        return output.getvalue()

    def save_json(self, filepath: str, indent: int = 2) -> None:
        """
        Save results to JSON file.

        Args:
            filepath: Path to output file
            indent: JSON indentation level
        """
        with open(filepath, 'w') as f:
            f.write(self.to_json(indent=indent))

    def save_csv(
        self,
        filepath: str,
        format: str = "long",
        include_metadata: bool = True
    ) -> None:
        """
        Save results to CSV file.

        Args:
            filepath: Path to output file
            format: "long" or "wide"
            include_metadata: Include temperature and other metadata
        """
        with open(filepath, 'w') as f:
            f.write(self.to_csv(format=format, include_metadata=include_metadata))

    def __len__(self) -> int:
        """Return number of data points"""
        return len(self._data_points)

    def __repr__(self) -> str:
        """String representation"""
        duration = ""
        if self.start_time and self.end_time:
            duration_sec = self.end_time - self.start_time
            duration = f", duration={duration_sec:.0f}s"

        return (f"ExperimentResults(id='{self.experiment_id}', "
                f"name='{self.experiment_name}', "
                f"measurements={len(self._data_points)}, "
                f"wells={len(self.get_wells())}{duration})")


class DataTransmissionParser:
    """
    Parser for CDCL data transmission messages.
    """

    @staticmethod
    def parse_data_message(
        message: str,
        results: ExperimentResults
    ) -> None:
        """
        Parse a data transmission message and add to results.

        Args:
            message: Raw CDCL message (typically starts with D prefix)
            results: ExperimentResults object to populate

        """
        response = CDCLParser.parse_response(message)
        if response.prefix == GenericResponsePrefixes.DATA_TRANSMISSION:
            experiment_id = response.params[0].strip().replace('\"', '')
            print(response.params)
            if experiment_id != results.experiment_id:
                print("Data transmission is for a different experiment")
                return results
            timestamp = int(response.params[1])

            if not isinstance(timestamp, int):
                print('timestamp not an integer')
            temperature = response.params[2]
            # mode is the 4th item
            mode = int(response.params[3])
            wavelength_index = PlateReaderModeCalculator.get_led_index(mode)
            wavelength = results.wavelengths[wavelength_index]
            for i in range(4, len(response.params)):
                well = index_to_well(i-4)
                print(f'well {well} value: {response.params[i]}')
                results.add_data_point(
                    timestamp=timestamp,
                    well=well,
                    value=response.params[i],
                    temperature=temperature,
                    wavelength=wavelength
                )

            return results
        else:
            raise ValueError(
                "Response is not a valid Data Transmission message")
