import numpy as np
from databroker.core import BlueskyRun
from xicam.core.data import ProjectionNotFound

from xicam.core.intents import ImageIntent

#TODO: check axes and transpose accrodingly

def project_nxSTXM(run_catalog: BlueskyRun):
    projection = next(filter(lambda projection: projection['name'] == 'nxSTXM', run_catalog.metadata['start']['projections']))['projection']
    stream = projection['irmap/DATA/data']['stream']
    field = projection['irmap/DATA/data']['field']
    sample_x = projection['irmap/DATA/sample_x']['value']
    sample_y = projection['irmap/DATA/sample_y']['value']
    energy = projection['irmap/DATA/energy']['value']

    xdata = getattr(run_catalog, stream).to_dask()[field]  # type: xr.DataArray

    xdata = np.squeeze(xdata)

    # xdata = xdata.assign_coords({xdata.dims[0]: energy, xdata.dims[1]: sample_y, xdata.dims[2]: sample_x})
    xdata = xdata.assign_coords({xdata.dims[0]: energy, xdata.dims[2]: sample_x, xdata.dims[1]: sample_y})

    return xdata.transpose('y (μm)', 'x (μm)', ...)


def project_nxCXI_ptycho(run_catalog: BlueskyRun):
    projection = next(
        filter(lambda projection: projection['name'] == 'NXcxi_ptycho', run_catalog.metadata['start']['projections']),
        None)
    if not projection:
        raise ProjectionNotFound

    transmission_rec_stream = projection['projection']['object_transmission']['stream']
    transmission_rec_field = projection['projection']['object_transmission']['field']
    phase_rec_stream = projection['projection']['object_phase']['stream']
    phase_rec_field = projection['projection']['object_phase']['field']
    energy = run_catalog.primary.read()['E [eV]'].data[0]
    coords_x = run_catalog.primary.read()['x [nm]'].data[0]
    coords_y = run_catalog.primary.read()['y [nm]'].data[0]

    # rec_data = getattr(run_catalog, transmission_rec_stream).to_dask().rename({transmission_rec_field: 'object_transmission', \
    #                                                                           phase_rec_field: 'object_phase'})
    rec_data_trans = getattr(run_catalog, transmission_rec_stream).to_dask()[transmission_rec_field]
    rec_data_trans = np.squeeze(rec_data_trans)
    rec_data_phase = getattr(run_catalog, phase_rec_stream).to_dask()[phase_rec_field]
    rec_data_phase = np.squeeze(rec_data_phase)

    rec_data_trans = rec_data_trans.assign_coords(
        {rec_data_trans.dims[0]: energy, rec_data_trans.dims[1]: coords_y, rec_data_trans.dims[2]: coords_x})

    return [ImageIntent(name='Transmission Reconstruction', image=rec_data_trans),
            # ImageIntent(image=rec_data_phase, item_name='phase reconstruction')
            ]


def project_arpes(catalog):
    # TODO: single-source
    ENERGY_FIELD = 'E (eV)'
    SAMPLE_X_FIELD = 'x (μm)'
    SAMPLE_Y_FIELD = 'y (μm)'
    ANGLE_FIELD = '???'

    data = catalog.primary.to_dask()
    raw_data = data['raw'][0]
    raw_data = raw_data.assign_coords({name: np.asarray(data[name][0])
                                       for name in raw_data.dims})
    return raw_data


projection_mapping = {'NXcxi_ptycho': project_nxCXI_ptycho,
                      'nxSTXM': project_nxSTXM,
                      'arpes': project_arpes}


def project_all(run_catalog: BlueskyRun):
    for projection in run_catalog.metadata['start']['projections']:
        projector = projection_mapping.get(projection['name'])
        if projector:
            return projector(run_catalog)
