# plate (repeated for N plates measured)
#     Plate information
#     Background information (optional)
#     Calculated results (optional)
#     Measured results (repeated for N measurements performed per plate)
# Basic assay information
# Protocol information
# Plate type
# Platemap (repeated for N plates measured)
# Calculations
# Auto export parameters
# Operations
# Labels
#     Filters (dependent on detection modality)
#     Mirrors (dependent on detection modality)
# Instrument
from __future__ import annotations

from collections import defaultdict
from collections.abc import Iterator
from dataclasses import dataclass
from enum import Enum
from re import search

import numpy as np

from allotropy.allotrope.models.adm.plate_reader.benchling._2023._09.plate_reader import (
    ScanPositionSettingPlateReader,
)
from allotropy.allotrope.models.shared.components.plate_reader import SampleRoleType
from allotropy.allotrope.models.shared.definitions.units import UNITLESS
from allotropy.allotrope.schema_mappers.adm.plate_reader.benchling._2023._09.plate_reader import (
    CalculatedDataItem,
    DataSource,
    Measurement,
    MeasurementGroup,
    MeasurementType,
    Metadata,
)
from allotropy.exceptions import AllotropeConversionError
from allotropy.parsers.constants import NOT_APPLICABLE
from allotropy.parsers.lines_reader import CsvReader
from allotropy.parsers.perkin_elmer_envision import constants
from allotropy.parsers.utils.pandas import (
    df_to_series_data,
    map_rows,
    SeriesData,
)
from allotropy.parsers.utils.uuids import random_uuid_str
from allotropy.parsers.utils.values import (
    assert_not_none,
    num_to_chars,
)


class ReadType(Enum):
    ABSORBANCE = "Absorbance"
    FLUORESCENCE = "Fluorescence"
    LUMINESCENCE = "Luminescence"

    @property
    def measurement_type(self) -> MeasurementType:
        if self is ReadType.ABSORBANCE:
            return MeasurementType.ULTRAVIOLET_ABSORBANCE
        elif self is ReadType.FLUORESCENCE:
            return MeasurementType.FLUORESCENCE
        elif self is ReadType.LUMINESCENCE:
            return MeasurementType.LUMINESCENCE

    @property
    def device_type(self) -> str:
        return f"{self.value.lower()} detector"


@dataclass(frozen=True)
class PlateInfo:
    number: str
    barcode: str
    measurement_time: str | None
    measured_height: float | None
    chamber_temperature_at_start: float | None

    @staticmethod
    def get_series(reader: CsvReader) -> SeriesData:
        assert_not_none(
            reader.pop_if_match("^Plate information"),
            msg="Unable to find expected plate information",
        )

        return df_to_series_data(
            assert_not_none(
                reader.pop_csv_block_as_df(header=0),
                "Plate information CSV block",
            ).replace(np.nan, None)
        )

    @staticmethod
    def get_plate_number_and_barcode(data: SeriesData) -> tuple[str, str]:
        plate_number = data[str, "Plate"]
        raw_barcode = data.get(str, "Barcode")
        barcode = (raw_barcode or '=""').removeprefix('="').removesuffix('"')
        barcode = barcode or f"Plate {plate_number}"
        return plate_number, barcode


@dataclass(frozen=True)
class CalculatedPlateInfo(PlateInfo):
    formula: str
    name: str

    @staticmethod
    def create(data: SeriesData) -> CalculatedPlateInfo:
        formula = data[str, "Formula"]

        name = assert_not_none(
            search(r"^([^=]*)=", formula),
            msg="Unable to find expected formula name for calculated results section.",
        ).group(1)

        plate_number, barcode = PlateInfo.get_plate_number_and_barcode(data)

        return CalculatedPlateInfo(
            number=plate_number,
            barcode=barcode,
            measurement_time=data.get(str, "Measurement date"),
            measured_height=data.get(float, "Measured height"),
            chamber_temperature_at_start=data.get(
                float, "Chamber temperature at start"
            ),
            formula=formula,
            name=name.strip(),
        )


@dataclass(frozen=True)
class ResultPlateInfo(PlateInfo):
    label: str
    measinfo: str
    emission_filter_id: str

    @staticmethod
    def create(data: SeriesData) -> ResultPlateInfo | None:
        label = data.get(str, "Label")
        if label is None:
            return None

        measinfo = data.get(str, "Measinfo")
        if measinfo is None:
            return None

        plate_number, barcode = PlateInfo.get_plate_number_and_barcode(data)

        return ResultPlateInfo(
            plate_number,
            barcode,
            data.get(str, "Measurement date"),
            data.get(float, "Measured height"),
            data.get(float, "Chamber temperature at start"),
            label=label,
            measinfo=measinfo,
            emission_filter_id=assert_not_none(
                search("De=(...)", measinfo),
                msg=f"Unable to find emission filter ID for Plate {barcode}.",
            ).group(1),
        )

    def match(self, background_info: BackgroundInfo) -> bool:
        return all(
            [
                background_info.plate_num == self.number,
                background_info.label in self.label,
                background_info.measinfo == self.measinfo,
            ]
        )


@dataclass
class BackgroundInfo:
    plate_num: str
    label: str
    measinfo: str

    @staticmethod
    def create(data: SeriesData) -> BackgroundInfo:
        return BackgroundInfo(
            plate_num=data[str, "Plate"],
            label=data[str, "Label"],
            measinfo=data[str, "MeasInfo"],
        )


def create_background_infos(reader: CsvReader) -> list[BackgroundInfo]:
    title = reader.pop_if_match("^Background information")
    if title is None:
        return []

    data_frame = assert_not_none(
        reader.pop_csv_block_as_df(header=0),
        "background information",
    )
    return map_rows(data_frame, BackgroundInfo.create)


@dataclass
class CalculatedResult:
    uuid: str
    col: str
    row: str
    value: float


def create_calculated_results(reader: CsvReader) -> list[CalculatedResult]:
    # Calculated results may or may not have a title
    reader.pop_if_match("^Calculated results")

    data_frame = assert_not_none(
        reader.pop_csv_block_as_df(),
        "results data",
    )
    series = (
        data_frame.drop(0, axis="index").drop(0, axis="columns")
        if data_frame.iloc[1, 0] == "A"
        else data_frame
    )
    rows, cols = series.shape
    series.index = [num_to_chars(i) for i in range(rows)]  # type: ignore[assignment]
    series.columns = [str(i).zfill(2) for i in range(1, cols + 1)]  # type: ignore[assignment]

    return [
        CalculatedResult(
            uuid=random_uuid_str(),
            col=col,
            row=row,
            value=series.loc[col, row],
        )
        for col, row in series.stack().index
    ]


@dataclass
class Result:
    uuid: str
    col: str
    row: str
    value: int


def create_results(reader: CsvReader) -> list[Result]:
    # Results may or may not have a title
    reader.pop_if_match("^Results")

    data_frame = assert_not_none(
        reader.pop_csv_block_as_df(),
        "reader data",
    )
    series = (
        data_frame.drop(0, axis="index").drop(0, axis="columns")
        if data_frame.iloc[1, 0] == "A"
        else data_frame
    )
    rows, cols = series.shape
    series.index = [num_to_chars(i) for i in range(rows)]  # type: ignore[assignment]
    series.columns = [str(i).zfill(2) for i in range(1, cols + 1)]  # type: ignore[assignment]

    return [
        Result(
            uuid=random_uuid_str(),
            col=col,
            row=row,
            value=int(series.loc[col, row]),
        )
        for col, row in series.stack().index
    ]


@dataclass
class Plate:
    plate_info: PlateInfo
    background_infos: list[BackgroundInfo]


@dataclass
class ResultPlate(Plate):
    plate_info: ResultPlateInfo
    results: list[Result]


@dataclass
class CalculatedPlate(Plate):
    plate_info: CalculatedPlateInfo
    results: list[CalculatedResult]

    def get_source_results(self, plate_list: PlateList) -> Iterator[list[Result]]:
        if not self.background_infos:
            msg = f"Unable to collect result plates for calculated data, there is no background information for plate {self.plate_info.number}"
            raise AllotropeConversionError(msg)

        source_results = [
            assert_not_none(
                plate_list.get_result_plate(background_info),
                msg=f"Unable to find result plate {background_info.label}.",
            ).results
            for background_info in self.background_infos
        ]

        if not all(len(source_results[0]) == len(sr) for sr in source_results):
            msg = f"Unable to collect result plates for calculated data, expected all result plates to have the same number of results, get: {[len(sr) for sr in source_results]}."
            raise AllotropeConversionError(msg)

        yield from [list(tup) for tup in zip(*source_results, strict=True)]

    def get_result_and_sources(
        self, plate_list: PlateList
    ) -> Iterator[tuple[CalculatedResult, list[Result]]]:
        yield from zip(self.results, self.get_source_results(plate_list), strict=True)


def create_plate(reader: CsvReader) -> ResultPlate | CalculatedPlate:
    series = PlateInfo.get_series(reader)
    if result_plate_info := ResultPlateInfo.create(series):
        return ResultPlate(
            plate_info=result_plate_info,
            background_infos=create_background_infos(reader),
            results=create_results(reader),
        )
    else:
        return CalculatedPlate(
            plate_info=CalculatedPlateInfo.create(series),
            background_infos=create_background_infos(reader),
            results=create_calculated_results(reader),
        )


@dataclass
class PlateList:
    results: list[ResultPlate]
    calculated: list[CalculatedPlate]

    @staticmethod
    def create(reader: CsvReader) -> PlateList:
        results: list[ResultPlate] = []
        calculated: list[CalculatedPlate] = []

        while reader.match("^Plate information"):
            plate = create_plate(reader)
            if isinstance(plate, CalculatedPlate):
                calculated.append(plate)
            else:
                results.append(plate)
        return PlateList(results, calculated)

    def get_result_plate(self, background_info: BackgroundInfo) -> ResultPlate | None:
        for plate in self.results:
            if plate.plate_info.match(background_info):
                return plate
        return None

    def get_measurement_time(self) -> str:
        try:
            return min(
                [
                    plate.plate_info.measurement_time
                    for plate in self.results
                    if plate.plate_info.measurement_time
                ]
            )
        except ValueError as err:
            msg = "Unable to determine the measurement time."
            raise AllotropeConversionError(msg) from err


@dataclass(frozen=True)
class BasicAssayInfo:
    protocol_id: str | None
    assay_id: str | None

    @staticmethod
    def create(reader: CsvReader) -> BasicAssayInfo:
        reader.drop_until_inclusive("^Basic assay information")
        data_frame = assert_not_none(
            reader.pop_csv_block_as_df(index_col=0),
            "Basic assay information",
        ).T
        data_frame.columns = data_frame.columns.str.replace(":", "").str.strip()
        data = df_to_series_data(data_frame)
        return BasicAssayInfo(data.get(str, "Protocol ID"), data.get(str, "Assay ID"))


@dataclass(frozen=True)
class PlateType:
    number_of_wells: float

    @staticmethod
    def create(reader: CsvReader) -> PlateType:
        reader.drop_until_inclusive("^Plate type")
        data_frame = assert_not_none(
            reader.pop_csv_block_as_df(index_col=0), "Plate type"
        )
        data = df_to_series_data(data_frame.T)
        return PlateType(
            number_of_wells=data[float, "Number of the wells in the plate"]
        )


def get_sample_role_type(encoding: str) -> SampleRoleType:
    # BL        blank               blank_role
    # CTL       control             control_sample_role
    # LB        lance_blank         blank_role
    # LC        lance_crosstalk     control_sample_role
    # LH        lance_high          control_sample_role
    # S         pl_sample           sample_role
    # STD       standard            standard_sample_role
    # -         unknown             unknown_sample_role
    # UNK       unknown             unknown_sample_role
    # ZH        z_high              control_sample_role
    # ZL        z_low               control_sample_role
    sample_role_type_map = {
        "BL": SampleRoleType.blank_role,
        "CTL": SampleRoleType.control_sample_role,
        "LB": SampleRoleType.blank_role,
        "LC": SampleRoleType.control_sample_role,
        "LH": SampleRoleType.control_sample_role,
        "STD": SampleRoleType.standard_sample_role,
        "S": SampleRoleType.sample_role,
        "-": SampleRoleType.unknown_sample_role,
        "UNK": SampleRoleType.unknown_sample_role,
        "ZH": SampleRoleType.control_sample_role,
        "ZL": SampleRoleType.control_sample_role,
    }
    for pattern, value in sample_role_type_map.items():
        if encoding.startswith(pattern):
            return value

    msg = f"Unable to determine sample role type of plate map encoding {encoding}; expected to start with one of {sorted(sample_role_type_map.keys())}."
    raise AllotropeConversionError(msg)


@dataclass(frozen=True)
class PlateMap:
    plate_n: str
    group_n: str
    sample_role_type_mapping: dict[str, dict[str, SampleRoleType]]

    @staticmethod
    def create(reader: CsvReader) -> PlateMap | None:
        if not reader.current_line_exists() or reader.match("^Calculations"):
            return None

        *_, plate_n = assert_not_none(reader.pop(), "Platemap number").split(",")
        *_, group_n = assert_not_none(reader.pop(), "Platemap group").split(",")

        data_frame = assert_not_none(
            reader.pop_csv_block_as_df(),
            "Platemap data",
        ).replace(" ", "", regex=True)

        reader.pop_data()  # drop type specification
        reader.drop_empty()

        series = (
            data_frame.drop(0, axis="index").drop(0, axis="columns")
            if data_frame.iloc[1, 0] == "A"
            else data_frame
        )
        rows, cols = series.shape
        series.index = [num_to_chars(i) for i in range(rows)]  # type: ignore[assignment]
        series.columns = [str(i).zfill(2) for i in range(1, cols + 1)]  # type: ignore[assignment]

        sample_role_type_mapping: dict[str, dict[str, SampleRoleType]] = {}
        for row, row_data in series.replace([np.nan, "''"], None).to_dict().items():
            col_mapping: dict[str, SampleRoleType] = {}
            for col, value in row_data.items():
                if value:
                    if role_type := get_sample_role_type(str(value)):
                        col_mapping[str(col)] = role_type
            if col_mapping:
                sample_role_type_mapping[str(row)] = col_mapping

        return PlateMap(plate_n, group_n, sample_role_type_mapping)

    def get_sample_role_type(self, col: str, row: str) -> SampleRoleType:
        try:
            return self.sample_role_type_mapping[row][col]
        except KeyError as e:
            msg = (
                f"Invalid plate map location for plate map {self.plate_n}: {col} {row}."
            )
            raise AllotropeConversionError(msg) from e


def create_plate_maps(reader: CsvReader) -> dict[str, PlateMap]:
    assert_not_none(
        reader.drop_until("^Platemap"),
        msg="No 'Platemap' section found.",
    )

    reader.pop()  # remove title

    maps: dict[str, PlateMap] = {}
    while _map := PlateMap.create(reader):
        maps[_map.plate_n] = _map
    return maps


@dataclass(frozen=True)
class Filter:
    name: str
    wavelength: float
    bandwidth: float | None = None

    @staticmethod
    def create(reader: CsvReader) -> Filter | None:
        if not reader.current_line_exists() or reader.match(
            "(^Mirror modules)|(^Instrument:)|(^Aperture:)"
        ):
            return None

        data_frame = assert_not_none(
            reader.pop_csv_block_as_df(index_col=0),
            "Filter information",
        )
        data = df_to_series_data(data_frame.T)
        name = str(data.series.index[0])
        description = data[str, "Description"]

        if search_result := search("Longpass=(\\d+)nm", description):
            return Filter(name, wavelength=float(search_result.group(1)))
        return Filter(
            name,
            wavelength=float(
                assert_not_none(
                    search("CWL=(\\d+)nm", description),
                    msg=f"Unable to find wavelength for filter {name}.",
                ).group(1)
            ),
            bandwidth=float(
                assert_not_none(
                    search("BW=(\\d+)nm", description),
                    msg=f"Unable to find bandwidth for filter {name}.",
                ).group(1)
            ),
        )


def create_filters(reader: CsvReader) -> dict[str, Filter]:
    reader.drop_until("(^Filters:)|^Instrument:")

    if reader.match("^Instrument"):
        return {}

    reader.pop()  # remove title

    filters = {}
    while _filter := Filter.create(reader):
        filters[_filter.name] = _filter
    return filters


@dataclass(frozen=True)
class Labels:
    label: str
    excitation_filter: Filter | None
    emission_filters: dict[str, Filter | None]
    scan_position_setting: ScanPositionSettingPlateReader | None = None
    number_of_flashes: float | None = None
    detector_gain_setting: str | None = None

    @staticmethod
    def create(reader: CsvReader) -> Labels:
        reader.drop_until_inclusive("^Labels")
        data_frame = assert_not_none(
            reader.pop_csv_block_as_df(index_col=0),
            "Labels",
        )
        data = df_to_series_data(data_frame.T.replace(np.nan, None))
        filters = create_filters(reader)
        filter_position_map = {
            "Bottom": ScanPositionSettingPlateReader.bottom_scan_position__plate_reader_,
            "Top": ScanPositionSettingPlateReader.top_scan_position__plate_reader_,
        }

        return Labels(
            label=data.series.index[0],
            excitation_filter=filters.get(data.get(str, "Exc. filter", NOT_APPLICABLE)),
            emission_filters={
                "1st": filters.get(
                    data.get(str, "Ems. filter", NOT_APPLICABLE),
                ),
                "2nd": filters.get(data.get(str, "2nd ems. filter", NOT_APPLICABLE)),
            },
            scan_position_setting=filter_position_map.get(
                data.get(str, "Using of emission filter", NOT_APPLICABLE)
            ),
            number_of_flashes=data.get(float, "Number of flashes"),
            detector_gain_setting=data.get(str, "Reference AD gain"),
        )

    def get_emission_filter(self, id_val: str) -> Filter | None:
        return self.emission_filters.get(id_val)

    def get_read_type(self) -> ReadType:
        patterns = {
            "ABS": ReadType.ABSORBANCE,
            "Absorbance": ReadType.ABSORBANCE,
            "LUM": ReadType.LUMINESCENCE,
            "Luminescence": ReadType.LUMINESCENCE,
            "Fluorescence": ReadType.FLUORESCENCE,
        }
        read_types = {
            read_type for key, read_type in patterns.items() if key in self.label
        }

        if len(read_types) > 1:
            msg = f"Unable to determine unique read type from label: '{self.label}'"
            raise AllotropeConversionError(msg)
        if len(read_types) == 1:
            return read_types.pop()
        # TODO check if this is correct, this is the original behavior
        return ReadType.FLUORESCENCE


@dataclass(frozen=True)
class Instrument:
    serial_number: str
    nickname: str

    @staticmethod
    def create(reader: CsvReader) -> Instrument:
        assert_not_none(
            reader.drop_until("^Instrument"),
            msg="No 'Instrument' section found.",
        )

        reader.pop()  # remove title

        *_, serial_number = assert_not_none(reader.pop(), "serial number").split(",")
        *_, nickname = assert_not_none(reader.pop(), "nickname").split(",")

        return Instrument(serial_number, nickname)


@dataclass(frozen=True)
class Software:
    software_name: str
    software_version: str

    @staticmethod
    def create(reader: CsvReader) -> Software:
        exported_with_regex = "Exported with (.+) version (.+)"
        assert_not_none(
            reader.drop_until(exported_with_regex),
            msg="Unable to find software information; no 'Exported with' section found.",
        )

        search_result = assert_not_none(
            search(
                exported_with_regex,
                assert_not_none(
                    reader.pop(),
                    "software information",
                ),
            )
        )

        return Software(
            software_name=search_result.group(1),
            software_version=search_result.group(2),
        )


@dataclass(frozen=True)
class Data:
    software: Software
    plate_list: PlateList
    basic_assay_info: BasicAssayInfo
    number_of_wells: float
    plate_maps: dict[str, PlateMap]
    labels: Labels
    instrument: Instrument

    @staticmethod
    def create(reader: CsvReader) -> Data:
        return Data(
            plate_list=PlateList.create(reader),
            basic_assay_info=BasicAssayInfo.create(reader),
            number_of_wells=PlateType.create(reader).number_of_wells,
            plate_maps=create_plate_maps(reader),
            labels=Labels.create(reader),
            instrument=Instrument.create(reader),
            software=Software.create(reader),
        )


def create_metadata(
    software: Software, instrument: Instrument, file_name: str
) -> Metadata:
    return Metadata(
        file_name=file_name,
        software_name=software.software_name,
        software_version=software.software_version,
        model_number=constants.MODEL_NUMBER,
        equipment_serial_number=instrument.serial_number,
        device_identifier=instrument.nickname,
    )


def _create_measurement(
    plate_info: ResultPlateInfo,
    result: Result,
    plate_maps: dict[str, PlateMap],
    labels: Labels,
    read_type: ReadType,
) -> Measurement:
    try:
        p_map = plate_maps[plate_info.number]
    except KeyError as e:
        msg = f"Unable to find plate map for Plate {plate_info.barcode}."
        raise AllotropeConversionError(msg) from e
    plate_barcode = plate_info.barcode
    well_location = f"{result.col}{result.row}"
    ex_filter = labels.excitation_filter
    em_filter = labels.get_emission_filter(plate_info.emission_filter_id)
    return Measurement(
        type_=read_type.measurement_type,
        device_type=read_type.device_type,
        identifier=result.uuid,
        sample_identifier=f"{plate_barcode} {well_location}",
        well_plate_identifier=plate_barcode,
        location_identifier=well_location,
        sample_role_type=p_map.get_sample_role_type(result.col, result.row),
        compartment_temperature=plate_info.chamber_temperature_at_start,
        absorbance=result.value if read_type is ReadType.ABSORBANCE else None,
        fluorescence=result.value if read_type is ReadType.FLUORESCENCE else None,
        luminescence=result.value if read_type is ReadType.LUMINESCENCE else None,
        detector_distance_setting=plate_info.measured_height,
        number_of_averages=labels.number_of_flashes,
        detector_gain_setting=labels.detector_gain_setting,
        scan_position_setting=labels.scan_position_setting,
        detector_wavelength_setting=em_filter.wavelength if em_filter else None,
        detector_bandwidth_setting=em_filter.bandwidth if em_filter else None,
        excitation_wavelength_setting=ex_filter.wavelength if ex_filter else None,
        excitation_bandwidth_setting=ex_filter.bandwidth if ex_filter else None,
    )


def create_measurement_groups(data: Data) -> list[MeasurementGroup]:
    read_type = data.labels.get_read_type()
    well_loc_measurements = defaultdict(list)
    for plate in data.plate_list.results:
        for result in plate.results:
            measurement = _create_measurement(
                plate.plate_info,
                result,
                data.plate_maps,
                data.labels,
                read_type,
            )
            well_loc_measurements[
                (plate.plate_info.number, measurement.location_identifier)
            ].append(measurement)

    measurement_time = data.plate_list.get_measurement_time()
    return [
        MeasurementGroup(
            measurement_time=measurement_time,
            plate_well_count=data.number_of_wells,
            analytical_method_identifier=data.basic_assay_info.protocol_id,
            experimental_data_identifier=data.basic_assay_info.assay_id,
            measurements=well_loc_measurements[well_location],
        )
        for well_location in sorted(well_loc_measurements.keys())
    ]


def create_calculated_data(
    plate_list: PlateList, read_type: ReadType
) -> list[CalculatedDataItem]:
    return [
        CalculatedDataItem(
            identifier=calculated_result.uuid,
            name=calculated_plate.plate_info.name,
            description=calculated_plate.plate_info.formula,
            value=calculated_result.value,
            unit=UNITLESS,
            data_sources=[
                DataSource(
                    identifier=source_result.uuid,
                    feature=read_type.value,
                )
                for source_result in source_results
            ],
        )
        for calculated_plate in plate_list.calculated
        for calculated_result, source_results in calculated_plate.get_result_and_sources(
            plate_list
        )
    ]
