from datetime import datetime
from unittest import TestCase
from uuid import uuid4

import numpy as np
from pynwb import NWBFile
from pynwb.ophys import OpticalChannel, ImageSegmentation, RoiResponseSeries
from hdmf.common.table import DynamicTableRegion, DynamicTable

from nwbinspector.checks.ophys import (
    check_roi_response_series_dims,
    check_roi_response_series_link_to_plane_segmentation,
)
from nwbinspector.register_checks import InspectorMessage, Importance


class TestCheckRoiResponseSeries(TestCase):
    def setUp(self):

        nwbfile = NWBFile(
            session_description="", identifier=str(uuid4()), session_start_time=datetime.now().astimezone()
        )

        device = nwbfile.create_device(
            name="Microscope", description="My two-photon microscope", manufacturer="The best microscope manufacturer"
        )
        optical_channel = OpticalChannel(name="OpticalChannel", description="an optical channel", emission_lambda=500.0)
        imaging_plane = nwbfile.create_imaging_plane(
            name="ImagingPlane",
            optical_channel=optical_channel,
            imaging_rate=30.0,
            description="a very interesting part of the brain",
            device=device,
            excitation_lambda=600.0,
            indicator="GFP",
            location="V1",
            grid_spacing=[0.01, 0.01],
            grid_spacing_unit="meters",
            origin_coords=[1.0, 2.0, 3.0],
            origin_coords_unit="meters",
        )

        img_seg = ImageSegmentation()

        self.plane_segmentation = img_seg.create_plane_segmentation(
            name="PlaneSegmentation",
            description="output from segmenting my favorite imaging plane",
            imaging_plane=imaging_plane,
        )

        self.ophys_module = nwbfile.create_processing_module(
            name="ophys",
            description="optical physiology processed data",
        )

        self.ophys_module.add(img_seg)

        for _ in range(10):
            image_mask = np.zeros((100, 100))
            self.plane_segmentation.add_roi(image_mask=image_mask)
        self.nwbfile = nwbfile

    def test_check_flipped_dims(self):

        rt_region = self.plane_segmentation.create_roi_table_region(
            region=[0, 1, 2, 3, 4],
            description="the first of two ROIs",
        )

        roi_resp_series = RoiResponseSeries(
            name="RoiResponseSeries",
            data=np.ones((5, 40)),  # 50 samples, 2 ROIs
            rois=rt_region,
            unit="n.a.",
            rate=30.0,
        )

        self.ophys_module.add(roi_resp_series)

        assert check_roi_response_series_dims(roi_resp_series) == InspectorMessage(
            message="The second dimension of data does not match the length of rois, "
            "but instead the first does. Data is oriented incorrectly and should be transposed.",
            importance=Importance.CRITICAL,
            check_function_name="check_roi_response_series_dims",
            object_type="RoiResponseSeries",
            object_name="RoiResponseSeries",
        )

    def test_check_wrong_dims(self):

        rt_region = self.plane_segmentation.create_roi_table_region(
            region=[0, 1, 2, 3, 4],
            description="the first of two ROIs",
        )

        roi_resp_series = RoiResponseSeries(
            name="RoiResponseSeries",
            data=np.ones((10, 40)),  # 50 samples, 2 ROIs
            rois=rt_region,
            unit="n.a.",
            rate=30.0,
        )

        self.ophys_module.add(roi_resp_series)

        assert check_roi_response_series_dims(roi_resp_series) == InspectorMessage(
            message="The second dimension of data does not match the length of rois. Your " "data may be transposed.",
            importance=Importance.CRITICAL,
            check_function_name="check_roi_response_series_dims",
            object_type="RoiResponseSeries",
            object_name="RoiResponseSeries",
        )

    def test_pass_check_roi_response_series_dims(self):
        rt_region = self.plane_segmentation.create_roi_table_region(
            region=[0, 1, 2, 3, 4],
            description="the first of two ROIs",
        )

        roi_resp_series = RoiResponseSeries(
            name="RoiResponseSeries",
            data=np.ones((40, 5)),  # 50 samples, 2 ROIs
            rois=rt_region,
            unit="n.a.",
            rate=30.0,
        )

        assert check_roi_response_series_dims(roi_resp_series) is None

    def test_check_roi_response_series_link_to_plane_segmentation(self):
        dt = DynamicTable(name="name", description="desc")
        dt.add_column("a", "desc")
        for _ in range(5):
            dt.add_row(a=1)
        dtr = DynamicTableRegion(name="n", description="desc", data=[0, 1, 2, 3, 4], table=dt)
        roi_resp_series = RoiResponseSeries(
            name="RoiResponseSeries",
            data=np.ones((40, 5)),  # 50 samples, 2 ROIs
            rois=dtr,
            unit="n.a.",
            rate=30.0,
        )

        self.ophys_module.add(roi_resp_series)

        assert check_roi_response_series_link_to_plane_segmentation(roi_resp_series) == InspectorMessage(
            message="rois field does not point to a PlaneSegmentation table.",
            importance=Importance.BEST_PRACTICE_VIOLATION,
            check_function_name="check_roi_response_series_link_to_plane_segmentation",
            object_type="RoiResponseSeries",
            object_name="RoiResponseSeries",
        )

    def test_pass_check_roi_response_series_link_to_plane_segmentation(self):

        rt_region = self.plane_segmentation.create_roi_table_region(
            region=[0, 1, 2, 3, 4],
            description="the first of two ROIs",
        )

        roi_resp_series = RoiResponseSeries(
            name="RoiResponseSeries",
            data=np.ones((40, 5)),  # 50 samples, 2 ROIs
            rois=rt_region,
            unit="n.a.",
            rate=30.0,
        )

        assert check_roi_response_series_link_to_plane_segmentation(roi_resp_series) is None
