"""Functions that select points from a point cloud and put them in a new point cloud."""

import math
import os
import shapefile
import shapely
import sys
from shapely.geometry import Point
from shapely.errors import WKTReadingError
from shapely.wkt import loads
from shapely.geometry import box
from shapely.vectorized import contains
import numpy as np

from laserchicken.keys import point
from laserchicken import kd_tree
from laserchicken.utils import copy_point_cloud, add_metadata


def select_equal(point_cloud, attribute, value, return_mask=False):
    """
    Return the selection of the input point cloud that contains only points with a given attribute equal to some value.
    If a list of values is given, select the points corresponding to any of the provided values.

    :param point_cloud: Input point cloud.
    :param attribute: The attribute name used for selection
    :param value: The value(s) to compare the attribute to
    :param return_mask: If true, return the mask corresponding to the selection
    :return:
    """
    _check_valid_arguments(attribute, point_cloud)
    # broadcast using shape of the values
    mask = point_cloud[point][attribute]['data'] == np.array(value)[..., None]
    if mask.ndim > 1:
        mask = np.any(mask, axis=0)  # reduce
    if return_mask:
        return mask
    point_cloud_filtered = copy_point_cloud(point_cloud, mask)
    add_metadata(point_cloud_filtered, sys.modules[__name__],
                 {'attribute': attribute, 'value': value})
    return point_cloud_filtered


def select_above(point_cloud, attribute, threshold, return_mask=False):
    """
    Return the selection of the input point cloud that contains only points with a given attribute above some value.

    :param point_cloud: Input point cloud
    :param attribute: The attribute name used for selection
    :param threshold: The threshold value used for selection
    :param return_mask: If true, return the mask corresponding to the selection
    :return:
    """
    _check_valid_arguments(attribute, point_cloud)
    mask = point_cloud[point][attribute]['data'] > threshold
    if return_mask:
        return mask
    point_cloud_filtered = copy_point_cloud(point_cloud, mask)
    add_metadata(point_cloud_filtered, sys.modules[__name__],
                 {'attribute': attribute, 'threshold': threshold})
    return point_cloud_filtered


def select_below(point_cloud, attribute, threshold, return_mask=False):
    """
    Return the selection of the input point cloud that contains only points with a given attribute below some value.

    :param point_cloud: Input point cloud
    :param attribute: The attribute name used for selection
    :param threshold: The threshold value used for selection
    :param return_mask: If true, return the mask corresponding to the selection
    :return:
    """
    _check_valid_arguments(attribute, point_cloud)
    mask = point_cloud[point][attribute]['data'] < threshold
    if return_mask:
        return mask
    point_cloud_filtered = copy_point_cloud(point_cloud, mask)
    add_metadata(point_cloud_filtered, sys.modules[__name__],
                 {'attribute': attribute, 'threshold': threshold})
    return point_cloud_filtered


def _check_valid_arguments(attribute, point_cloud):
    """
    Raise if arguments are not valid for select_above/select_below functions.

    :param attribute:
    :param point_cloud:
    :return: None
    """
    if point_cloud is None:
        raise ValueError('Input point cloud cannot be None.')
    if attribute not in point_cloud[point]:
        raise ValueError('Attribute key {} for selection not found in point cloud.'.format(attribute))


def select_polygon(point_cloud, polygon_string, read_from_file=False, return_mask=False):
    """
    Return the selection of the input point cloud that contains only points within a given polygon.

    :param point_cloud: Input point cloud
    :param polygon_string: Polygon, either defined in a WKT string or in a file (WKT and ESRI formats supported)
    :param read_from_file: if true, polygon is expected to be the name of the file where the polygon is defined
    :param return_mask: if true, return a mask of selected points, rather than point cloud
    :return:
    """
    if point_cloud is None:
        raise ValueError('Input point cloud cannot be None.')
    if not isinstance(polygon_string, str):
        raise ValueError('Polygon (or its filename) should be a string')
    if read_from_file:
        format = os.path.splitext(polygon_string)[1].lower()
        reader = _get_polygon_reader(format)
        polygon = reader(polygon_string)
    else:
        polygon = _load_polygon(polygon_string)
    
    if isinstance(polygon, shapely.geometry.polygon.Polygon):
        points_in = _contains(point_cloud, polygon)
    elif isinstance(polygon,shapely.geometry.multipolygon.MultiPolygon):
        points_in = []
        count=1
        for poly in polygon:
            if not(count%200) or count==len(polygon):
                print('Checking polygon {}/{}...'.format(count, len(polygon)))
            points_in.extend(_contains(point_cloud, poly))
            count=count+1
        print('{} points found in {} polygons.'.format(len(points_in), len(polygon)))
    else:
        raise ValueError('It is not a Polygon or Multipolygon.')
    
    if return_mask: 
        mask = np.zeros(len(point_cloud['vertex']['x']['data']), dtype=bool)
        mask[points_in] = True
        return mask
    else:
        point_cloud_filtered = copy_point_cloud(point_cloud, points_in)
        add_metadata(point_cloud_filtered, sys.modules[__name__],
                    {'polygon_string': polygon_string,
                    'read_from_file': read_from_file})
        return point_cloud_filtered


def _read_wkt_file(path):
    with open(path) as f:
        content = f.readlines()

    content = [_load_polygon(x.strip()) for x in content]
    geom = shapely.geometry.MultiPolygon(content) if len(content) > 1 else content[0]
    return geom


def _read_shp_file(path):
    shape = shapefile.Reader(path)
    features = shape.shapeRecords()
    shp_geoms = [shapely.geometry.shape(feature.shape.__geo_interface__)
                 for feature in features]
    shp_geom = shapely.geometry.MultiPolygon(shp_geoms) if len(shp_geoms) > 1 else shp_geoms[0]
    return shp_geom


polygon_readers = {
    '.wkt': _read_wkt_file,
    '.shp': _read_shp_file
}


def _get_polygon_reader(format):
    if format not in polygon_readers: