from __future__ import print_function
import os
import warnings

import numpy as np
import rasterio

from rpcm import rpc_model


warnings.filterwarnings("ignore",
                        category=rasterio.errors.NotGeoreferencedWarning)


def viewing_direction(zenith, azimut):
    """
    Compute the unit 3D vector defined by zenith and azimut angles.

    Args:
        zenith (float): angle wrt the vertical direction, in degrees
        azimut (float): angle wrt the north direction, in degrees

    Return:
        3-tuple: 3D unit vector giving the corresponding direction
    """
    z = np.radians(zenith)
    a = np.radians(azimut)
    return np.sin(a)*np.sin(z), np.cos(a)*np.sin(z), np.cos(z)


def rpc_from_geotiff(geotiff_path):
    """
    Read the RPC coefficients from a GeoTIFF file and return a rpc_model object.

    Args:
        geotiff_path (str): path or url to a GeoTIFF file

    Returns:
        instance of the rpc_model.RPCModel class
    """
    with rasterio.open(geotiff_path, 'r') as src:
        rpc_dict = src.tags(ns='RPC')
    return rpc_model.RPCModel(rpc_dict)


def bounding_box2D(pts):
    """
    Rectangular bounding box for a list of 2D points.

    Args:
        pts (list): list of 2D points represented as 2-tuples or lists of length 2

    Returns:
        x, y, w, h (floats): coordinates of the top-left corner, width and
            height of the bounding box
    """
    dim = len(pts[0])  # should be 2
    bb_min = [min([t[i] for t in pts]) for i in range(dim)]
    bb_max = [max([t[i] for t in pts]) for i in range(dim)]
    return bb_min[0], bb_min[1], bb_max[0] - bb_min[0], bb_max[1] - bb_min[1]


def points_apply_homography(H, pts):
    """
    Applies an homography to a list of 2D points.

    Args:
        H (np.array): 3x3 homography matrix
        pts (list): list of 2D points, each point being a 2-tuple or a list
            with its x, y coordinates

    Returns:
        numpy array: list of transformed points, one per line
    """
    pts = np.asarray(pts)

    # convert the input points to homogeneous coordinates
    if len(pts[0]) < 2:
        print("""points_apply_homography: ERROR the input must be a numpy array
          of 2D points, one point per line""")
        return
    pts = np.hstack((pts[:, :2], np.ones(len(pts))))

    # apply the transformation
    Hpts = (np.dot(H, pts.T)).T

    # normalize the homogeneous result and trim the extra dimension
    Hpts = Hpts * (1.0 / np.tile( Hpts[:, 2], (3, 1)) ).T
    return Hpts[:, :2]


def bounding_box_of_projected_aoi(rpc, aoi, z=0, homography=None):
    """
    Return the x, y, w, h pixel bounding box of a projected AOI.

    Args:
        rpc (rpc_model.RPCModel): RPC camera model
        aoi (geojson.Polygon): GeoJSON polygon representing the AOI
        z (float): altitude of the AOI with respect to the WGS84 ellipsoid
        homography (2D array, optional): matrix of shape (3, 3) representing an
            homography to be applied to the projected points before computing
            their bounding box.

    Return:
        x, y (ints): pixel coordinates of the top-left corner of the bounding box
        w, h (ints): pixel dimensions of the bounding box
    """
    lons, lats = np.asarray(aoi['coordinates']).squeeze().T
    x, y = rpc.projection(lons, lats, z)
    pts = list(zip(x, y))
    if homography is not None:
        pts = points_apply_homography(homography, pts)
    return np.round(bounding_box2D(pts)).astype(int)


def crop_aoi(geotiff, aoi, z=0):
    """
    Crop a geographic AOI in a georeferenced image using its RPC functions.

    Args:
        geotiff (string): path or url to the input GeoTIFF image file
        aoi (geojson.Polygon): GeoJSON polygon representing the AOI
        z (float, optional): base altitude with respect to WGS84 ellipsoid (0
            by default)

    Return:
        crop (array): numpy array containing the cropped image
        x, y, w, h (ints): image coordinates of the crop. x, y are the
            coordinates of the top-left corner, while w, h are the dimensions
            of the crop.
    """
    x, y, w, h = bounding_box_of_projected_aoi(rpc_from_geotiff(geotiff), aoi, z)
    with rasterio.open(geotiff, 'r') as src:
        crop = src.read(window=((y, y + h), (x, x + w)), boundless=True).squeeze()
    return crop, x, y


def rasterio_write(path, array, profile={}, tags={}):
    """
    Write a numpy array in a tiff or png file with rasterio.

    Args:
        path (str): path to the output tiff/png file
        array (numpy array): 2D or 3D array containing the image to write.
        profile (dict): rasterio profile (ie dictionary of metadata)
        tags (dict): dictionary with additional geotiff tags
    """
    # determine the driver based on the file extension
    extension = os.path.splitext(path)[1].lower()
    if extension in ['.tif', '.tiff']:
        driver = 'GTiff'
    elif extension in ['.png']:
        driver = 'png'
    else:
        raise NotImplementedError('format {} not supported'.format(extension))

    # read image size and number of bands
    array = np.atleast_3d(array)
    height, width, nbands = array.shape

    # define image metadata dict
    profile.update(driver=driver, count=nbands, width=width, height=height,
                   dtype=array.dtype)

    # write to file
    with rasterio.Env():
        with rasterio.open(path, 'w', **profile) as dst:
            dst.write(np.transpose(array, (2, 0, 1)))
            dst.update_tags(**tags)


def compute_epsg(lon, lat):
    """
    Compute the EPSG code of the UTM zone which contains
    the point with given longitude and latitude

    Args:
        lon (float): longitude of the point
        lat (float): latitude of the point

    Returns:
        int: EPSG code
    """
    # UTM zone number starts from 1 at longitude -180,
    # and increments by 1 every 6 degrees of longitude
    zone = int((lon + 180) // 6 + 1)

    # EPSG = CONST + ZONE where CONST is
    # - 32600 for positive latitudes
    # - 32700 for negative latitudes
    const = 32600 if lat > 0 else 32700
    epsg = const + zone
    return epsg


def approximate_rpc_as_projective(rpc_model, col_range, lin_range, alt_range,
        verbose=False):
    """
    Return a projective camera matrix approximating the RPC projection function.
    
    The approximation is optimized on a sampling of the 3D region
    defined by the altitudes in alt_range and the image tile defined by
    col_range and lin_range.
    """
    ### step 1: generate cartesian coordinates of 3d points used to fit the
    ###         best projection matrix
    # get mesh points and convert them to geodetic then to geocentric
    # coordinates
    cols, lins, alts = generate_point_mesh(col_range, lin_range, alt_range)
    lons, lats, alts = rpc_model.direct_estimate(cols, lins, alts)
    x, y, z = geographiclib.geodetic_to_geocentric(lats, lons, alts)

    ### step 2: estimate the camera projection matrix from corresponding
    # 3-space and image entities
    world_points = np.vstack([x, y, z]).T
    image_points = np.vstack([cols, lins]).T
    P = estimation.camera_matrix(world_points, image_points)

    ### step 3: for debug, test the approximation error
    if verbose:
        # compute the projection error made by the computed matrix P, on the
        # used learning points
        colPROJ = np.zeros(len(x))
        linPROJ = np.zeros(len(x))
        for i in range(len(x)):
            v = np.dot(P, [[x[i]],[y[i]],[z[i]],[1]])
            colPROJ[i] = v[0]/v[2]
            linPROJ[i] = v[1]/v[2]

        print('approximate_rpc_as_projective: (min, max, mean)')
        print_distance_between_vectors(cols, colPROJ, 'cols')
        print_distance_between_vectors(lins, linPROJ, 'rows')

    return P


