import numpy as np
import xarray as xr

from ..utils.jaggedarray import flatten_jagged_array


MESH_ATTRS = {
    'cf_role': 'mesh_topology',
    'long_name': 'Topology data of 2D unstructured mesh',
    'topology_dimension': 2,
    'node_coordinates': 'x_of_node y_of_node',
    'face_node_connectivity': 'nodes_at_patch',
    'face_dimension': 'patch',
    'face_edge_connectivity': 'links_at_patch',
    'edge_node_connectivity': 'nodes_at_link',
    'edge_dimension': 'link',
}

DUAL_MESH_ATTRS = {
    'cf_role': 'mesh_topology',
    'long_name': 'Topology data of 2D unstructured mesh',
    'topology_dimension': 2,
    'node_coordinates': 'x_of_corner y_of_corner',
    'face_node_connectivity': 'corners_at_cell',
    'face_dimension': 'cell',
    'face_edge_connectivity': 'faces_at_cell',
    'edge_node_connectivity': 'corners_at_face',
    'edge_dimension': 'face',
}


def ugrid_from_structured_quad(coords, shape=None):
    from .structured_quad.structured_quad import (setup_nodes_at_link,
                                                  setup_links_at_patch,
                                                  setup_node_coords_structured)
    node_y_and_x = setup_node_coords_structured(coords, shape=shape)
    shape = node_y_and_x[0].shape

    return ugrid_from_unstructured(
        node_y_and_x,
        setup_nodes_at_link(shape),
        setup_links_at_patch(shape)) 


def ugrid_from_rectilinear(coords):
    from .structured_quad.structured_quad import setup_node_coords_rectilinear
    node_y_and_x = setup_node_coords_rectilinear(coords)
    return ugrid_from_structured_quad(node_y_and_x)


def ugrid_from_uniform_rectilinear(shape, spacing=1., origin=0.):
    from .structured_quad.structured_quad import setup_node_coords
    node_y_and_x = setup_node_coords(shape, spacing, origin)
    return ugrid_from_structured_quad(node_y_and_x)


def load_ugrid(ugrid):
    if not isinstance(ugrid, xr.Dataset):
        raise AssertionError('not an instance of xarray.Dataset')

    ds = xr.Dataset({'mesh': xr.DataArray(data='a', attrs=MESH_ATTRS)})

    meta = ugrid.a.attrs

    node_coords = meta['node_coordinates'].split()
    ds.update({'x_of_node': ugrid[node_coords[0]],
               'y_of_node': ugrid[node_coords[1]], })
    ds = ds.rename({ds.x_of_node.dims[0]: 'node'})

    ds.update({'nodes_at_link': ugrid[meta['edge_node_connectivity']]})
    ds = ds.rename({ds.nodes_at_link.dims[0]: 'link',
                    ds.nodes_at_link.dims[1]: 'Two', })

    ds.update({'links_at_patch': ugrid[meta['face_edge_connectivity']]})
    ds = ds.rename({ds.links_at_patch.dims[0]: 'patch',
                    ds.links_at_patch.dims[1]: 'max_patch_links', })

    return ds


def ugrid_as_dataset(grid):
    if isinstance(grid, types.StringTypes):
        ugrid = xr.open_dataset(grid)
    elif isinstance(grid, Graph):
        ugrid = xr.ds
    elif isinstance(grid, dict):
        ugrid = xr.Dataset.from_dict(d)
    elif isinstance(grid, xr.Dataset):
        ugrid = grid
    else:
        raise ValueError(
            'unable to convert {t} to a Dataset'.format(t=type(grid)))

    if not validate_ugrid(ugrid):
        raise ValueError('unable to convert to a ugrid Dataset')
    else:
        return grid


def ugrid_as_dual(ugrid):
    rename = {
        'node': 'corner',
        'link': 'face',
        'patch': 'cell',
        'x_of_node': 'x_of_corner',
        'y_of_node': 'y_of_corner',
        'nodes_at_link': 'corners_at_face',
        'links_at_patch': 'faces_at_cell',
        'max_patch_links': 'max_cell_faces',
    }
    return ugrid.rename(rename)


def ugrid_from_unstructured(node_y_and_x, links=None, patches=None):
    ugrid = xr.Dataset({'mesh': xr.DataArray(data='a', attrs=MESH_ATTRS)})

    update_node_coords(ugrid, node_y_and_x)

    if links is not None:
        update_nodes_at_link(ugrid, links)

    if patches is not None and 'nodes_at_link' in ugrid:
        update_links_at_patch(ugrid, patches)

    return ugrid


def update_node_coords(ugrid, node_y_and_x):
    node_y, node_x = (np.asarray(node_y_and_x[0], dtype=float),
                      np.asarray(node_y_and_x[1], dtype=float))
    y_of_node = xr.DataArray(
        data=node_y.reshape((-1, )),
        coords={'node': np.arange(node_y.size)},
        dims=('node', ),
        attrs={ 'long_name': 'y', 'units': 'm',})
    x_of_node = xr.DataArray(
        data=node_x.reshape((-1, )),
        coords={'node': np.arange(node_x.size)},
        dims=('node', ),
        attrs={ 'long_name': 'x', 'units': 'm',})
    ugrid.update({'y_of_node': y_of_node, 'x_of_node': x_of_node})

    return ugrid


def update_nodes_at_link(ugrid, node_links):
    node_links = np.asarray(node_links, dtype=np.int).reshape((-1, 2))
    nodes_at_link = xr.DataArray(
        data=node_links,
        dims=('link', 'Two'),
        attrs={'cf_role': 'edge_node_connectivity',
               'long_name': 'nodes a link tail and head',
               'start_index': 0})
    ugrid.update({'nodes_at_link': nodes_at_link})


def update_links_at_patch(ugrid, patches):
    from .matrix.at_patch import links_at_patch

    if len(patches) > 0:
        patches = flatten_jagged_array(patches, dtype=int)
    patch_links = links_at_patch(patches)
    links_at_patch = xr.DataArray(
        data=patch_links,
        dims=('patch', 'max_patch_links'),
        attrs={'cf_role': 'face_edge_connectivity',
               'long_name': 'Maps every face to its edges',
               'start_index': 0}
    )
    ugrid.update({'links_at_patch': links_at_patch})


def update_nodes_at_patch(ugrid, nodes_at_patch):
    nodes_at_patch = xr.DataArray(
        data=nodes_at_patch,
        dims=('patch', 'max_patch_links'),
        attrs={'cf_role': 'patch_node_connectivity',
               'long_name': 'nodes defining patches',
               'start_index': 0})
    ugrid.update({'nodes_at_patch': nodes_at_patch})


def update_node_at_cell(ugrid, node_at_cell):
    node_at_cell = xr.DataArray(
        data=node_at_cell,
        dims=('cell', ),
        attrs={'cf_role': 'node_cell_connectivity',
               'long_name': 'node at cell',
               'start_index': 0})
    ugrid.update({'node_at_cell': node_at_cell})


def update_nodes_at_face(ugrid, nodes_at_face):
    nodes_at_face = xr.DataArray(
        data=nodes_at_face,
        dims=('face', 'Two'),
        attrs={'cf_role': 'node_face_connectivity',
               'long_name': 'nodes at face',
               'start_index': 0})
    ugrid.update({'nodes_at_face': nodes_at_face})
