from pynwb.testing.mock.file import mock_NWBFile
from pynwb.testing.mock.ecephys import mock_ElectricalSeries
from ndx_franklab_novela import DataAcqDevice, Probe, Shank, ShanksElectrode, NwbElectrodeGroup
from pynwb import NWBHDF5IO
import numpy as np
from pathlib import Path


def main():
    nwbfile = mock_NWBFile(
        identifier="duplicate_electrode_id_bug_demo",
        session_description="Mock NWB file demonstrating duplicate electrode ID issue"
    )

    data_acq_device = DataAcqDevice(
        name="my_data_acq",
        system="my_system",
        amplifier="my_amplifier",
        adc_circuit="my_adc_circuit"
    )
    nwbfile.add_device(data_acq_device)

    probe1_shank1_electrodes = [
        ShanksElectrode(name="1", rel_x=0.0, rel_y=0.0, rel_z=0.0),
        ShanksElectrode(name="2", rel_x=0.0, rel_y=10.0, rel_z=0.0),
    ]
    probe1_shank2_electrodes = [
        ShanksElectrode(name="1", rel_x=0.0, rel_y=0.0, rel_z=0.0),
        ShanksElectrode(name="2", rel_x=0.0, rel_y=10.0, rel_z=0.0),
    ]
    
    probe1 = Probe(
        name="probe1",
        id=1,
        probe_type="tetrode",
        units="um",
        probe_description="first probe",
        contact_side_numbering=False,
        contact_size=10.0,
        shanks=[
            Shank(name="1", shanks_electrodes=probe1_shank1_electrodes),
            Shank(name="2", shanks_electrodes=probe1_shank2_electrodes),
        ],
    )
    nwbfile.add_device(probe1)

    probe2_shank1_electrodes = [
        ShanksElectrode(name="1", rel_x=0.0, rel_y=0.0, rel_z=0.0),
        ShanksElectrode(name="2", rel_x=0.0, rel_y=10.0, rel_z=0.0),
    ]
    probe2_shank2_electrodes = [
        ShanksElectrode(name="1", rel_x=0.0, rel_y=0.0, rel_z=0.0),
        ShanksElectrode(name="2", rel_x=0.0, rel_y=10.0, rel_z=0.0),
    ]
    
    probe2 = Probe(
        name="probe2",
        id=2,
        probe_type="tetrode",
        units="um",
        probe_description="second probe",
        contact_side_numbering=False,
        contact_size=10.0,
        shanks=[
            Shank(name="1", shanks_electrodes=probe2_shank1_electrodes),
            Shank(name="2", shanks_electrodes=probe2_shank2_electrodes),
        ],
    )
    nwbfile.add_device(probe2)

    electrode_group1 = NwbElectrodeGroup(
        name="electrode_group1",
        description="electrode group for probe1",
        location="CA1",
        device=probe1,
        targeted_location="CA1",
        targeted_x=0.0,
        targeted_y=0.0,
        targeted_z=0.0,
        units="um",
    )
    nwbfile.add_electrode_group(electrode_group1)

    electrode_group2 = NwbElectrodeGroup(
        name="electrode_group2",
        description="electrode group for probe2",
        location="CA1",
        device=probe2,
        targeted_location="CA1",
        targeted_x=0.0,
        targeted_y=0.0,
        targeted_z=0.0,
        units="um",
    )
    nwbfile.add_electrode_group(electrode_group2)

    extra_cols = ["probe_shank", "probe_electrode", "bad_channel", "ref_elect_id"]
    for col in extra_cols:
        nwbfile.add_electrode_column(name=col, description=f"description for {col}")

    electrode_counter = 0
    for probe_id, electrode_group in [(1, electrode_group1), (2, electrode_group2)]:
        for shank_id in [1, 2]:
            for electrode_id in [1, 2]:
                nwbfile.add_electrode(
                    location="CA1",
                    group=electrode_group,
                    probe_shank=shank_id,
                    probe_electrode=electrode_id,
                    bad_channel=False,
                    ref_elect_id=0,
                    x=0.0,
                    y=0.0,
                    z=0.0,
                )
                electrode_counter += 1

    electrodes = nwbfile.electrodes.create_region(
        name="electrodes",
        region=list(range(electrode_counter)),
        description="electrodes"
    )
    mock_ElectricalSeries(electrodes=electrodes, nwbfile=nwbfile, data=np.ones((10, 1)))

    nwbfile.create_processing_module(name="behavior", description="dummy behavior module")

    nwbfile_path = Path("/Volumes/T7/CatalystNeuro/Spyglass/raw/mock_duplicate_electrode_id_bug.nwb")
    if nwbfile_path.exists():
        nwbfile_path.unlink()
    nwbfile_path.parent.mkdir(parents=True, exist_ok=True)
    
    with NWBHDF5IO(nwbfile_path, "w") as io:
        io.write(nwbfile)

    import datajoint as dj
    dj_local_conf_path = "/Users/pauladkisson/Documents/CatalystNeuro/Spyglass/spyglass/dj_local_conf.json"
    dj.config.load(dj_local_conf_path)
    import spyglass.common as sgc
    import spyglass.data_import as sgi
    from spyglass.utils.nwb_helper_fn import get_nwb_copy_filename
    
    nwb_copy_file_name = get_nwb_copy_filename(nwbfile_path.name)
    (sgc.Nwbfile & {"nwb_file_name": nwb_copy_file_name}).delete()
    
    sgi.insert_sessions(str(nwbfile_path), rollback_on_fail=True, raise_err=True)


if __name__ == "__main__":
    main()