from plio.io.io_gdal import GeoDataset
import numpy as np
import os
import os.path

import pandas as pd

import geopandas as gpd

from geoalchemy2 import functions



from shapely import wkt
from shapely.geometry import Point

#from autocnet.matcher.subpixel import check_match_func
#from autocnet.matcher.subpixel import geom_match_simple
from autocnet.io.db.model import Images, Points, Measures, JsonEncoder
from autocnet.cg.cg import distribute_points_in_geom, xy_in_polygon
from autocnet.spatial import isis
from autocnet.transformation.spatial import reproject, oc2og
from autocnet.matcher.cpu_extractor import extract_most_interesting
from autocnet.transformation import roi
from autocnet.utils.utils import bytescale

import logging

log = logging.getLogger(__name__)

def generate_ground_points(Session, ground_mosaic, nspts_func=lambda x: int(round(x,1)*1), ewpts_func=lambda x: int(round(x,1)*4), size=(100,100)):
    """

    Parameters
    ----------
    ground_db_config : dict
                       In the form: {'username':'somename',
                                     'password':'somepassword',
                                     'host':'somehost',
                                     'pgbouncer_port':6543,
                                     'name':'somename'}
    nspts_func       : func
                       describes distribution of points along the north-south
                       edge of an overlap.

    ewpts_func       : func
                       describes distribution of points along the east-west
                       edge of an overlap.

    size             : tuple of int
                       (size_x, size_y) maximum distances on either access point
                       can move when attempting to find an interesting feature.
    """

    if isinstance(ground_mosaic, str):
        ground_mosaic = GeoDataset(ground_mosaic)

    log.warning('This function is not well tested. No tests currently exist \
    in the test suite for this version of the function.')

    session = Session()
    fp_poly = wkt.loads(session.query(functions.ST_AsText(functions.ST_Union(Images.geom))).one()[0])
    session.close()

    coords = distribute_points_in_geom(fp_poly, nspts_func=nspts_func, ewpts_func=ewpts_func, method="new", Session=Session)
    coords = np.asarray(coords)

    old_coord_list = []
    coord_list = []
    lines = []
    samples = []
    newlines = []
    newsamples = []

    # throw out points not intersecting the ground reference images
    log.info('points to lay down: ', len(coords))
    for i, coord in enumerate(coords):
        # res = ground_session.execute(formated_sql)
        p = Point(*coord)
        log.info(f'point {i}'),


        linessamples = isis.point_info(ground_mosaic.file_name, p.x, p.y, 'ground')
        if linessamples is None:
            log.warning('unable to find point in ground image')
            continue
        line = linessamples.get('Line')
        sample = linessamples.get('Sample')

        oldpoint = isis.point_info(ground_mosaic.file_name, sample, line, 'image')
        op = Point(oldpoint.get('PositiveEast360Longitude'),
                   oldpoint.get('PlanetocentricLatitude'))


        image = roi.Roi(ground_mosaic, sample, line, size_x=size[0], size_y=size[1])
        image_roi = image.clip(dtype="uint64")

        interesting = extract_most_interesting(bytescale(image_roi),  extractor_parameters={'nfeatures':30})

        # kps are in the image space with upper left origin, so convert to
        # center origin and then convert back into full image space
        left_x, _, top_y, _ = image.image_extent
        newsample = left_x + interesting.x
        newline = top_y + interesting.y

        newpoint = isis.point_info(ground_mosaic.file_name, newsample, newline, 'image')
        p = Point(newpoint.get('PositiveEast360Longitude'),
                  newpoint.get('PlanetocentricLatitude'))

        if not (xy_in_polygon(p.x, p.y, fp_poly)):
                log.warning('Interesting point not in mosaic area, ignore')
                continue

        old_coord_list.append(op)
        lines.append(line)
        samples.append(sample)
        coord_list.append(p)
        newlines.append(newline)
        newsamples.append(newsample)


    # start building the cnet
    ground_cnet = pd.DataFrame()
    ground_cnet["path"] = [ground_mosaic.file_name]*len(coord_list)
    ground_cnet["pointid"] = list(range(len(coord_list)))
    ground_cnet["original point"] = old_coord_list
    ground_cnet["point"] = coord_list
    ground_cnet['original_line'] = lines
    ground_cnet['original_sample'] = samples
    ground_cnet['line'] = newlines
    ground_cnet['sample'] = newsamples
    ground_cnet = gpd.GeoDataFrame(ground_cnet, geometry='point')
    return ground_cnet, fp_poly, coord_list


def propagate_point(Session,
                    config,
                    dem,
                    lon,
                    lat,
                    pointid,
                    paths,
                    lines,
                    samples,
                    size_x=40,
                    size_y=40,
                    match_func="classic",
                    match_kwargs={'image_size': (39, 39), 'template_size': (21, 21)},
                    verbose=False,
                    cost=lambda x, y: y == np.max(x)):
    """
    Conditionally propagate a point into a stack of images. The point and all corresponding measures
    are matched against database network (to which you are propagating), best result(s) is(are) kept.

    Parameters
    ----------
    Session : sqlalchemy.sessionmaker
              session maker associated with the database you want to propagate to

    config : dict
             configuration file associated with database you want to propagate to
             In the form: {'username':'somename', 'password':'somepassword',
             'host':'somehost', 'pgbouncer_port':6543, 'name':'somename'}

    dem : surface
          surface model of target body

    lon : np.float
          longitude of point you want to project

    lat : np.float
          planetocentric latitude of point you want to project

    pointid : int
             clerical input used to trace point from generate_ground_points output

    paths : list of str
            absolute paths pointing to the image(s) from which you want to try propagating the point 
    
    lines : list of np.float
            a priori line(s) corresponding to point projected in 'paths' image(s 

    samples : list of np.float
              a priori sample(s) corresponding to point projected in 'paths' image(s 

    size_x : int
             half width of GeoDataset that is cut from full image and affinely transfromed in geom_match;
             must be larger than 1/2 template_kwargs 'image_size 

    size_y : int
             half height of GeoDataset that is cut from full image and affinely transfromed in geom_match;
             must be larger than 1/2 template_kwargs 'image_size'

    template_kwargs : dict
                      kwargs passed through to control matcher.subpixel_template()

    verbose : boolean
              If True, this will print out the results of each propagation, including prints of the
              matcher areas and their correlation map.

    cost : anonymous function
           determines to which image(s) the point should be propagated. x corresponds to a list
           of all match correlation metrics, while y corresponds to each individual element
           of the x array.
           Example:
           cost = lambda x,y: y == np.max(x) will get you one result corresponding to the image that
           has the maximum correlation with the source image
           cost = lambda x,y: y > 0.6 will propagate the point to all images whose correlation
           result is greater than 0.6


    Returns
    -------
    new_measures : pd.DataFrame
                   Dataframe containing pointid, imageid, image serial number, line, sample, and ground location (both latlon
                   and cartesian) of successfully propagated points

    """

    match_func = check_match_func(match_func)

    session = Session()
    engine = session.get_bind()
    string = f"select * from images where ST_Intersects(geom, ST_SetSRID(ST_Point({lon}, {lat}), {config['spatial']['latitudinal_srid']}))"
    images = pd.read_sql(string, engine)
    session.close()

    image_measures = pd.DataFrame(zip(paths, lines, samples), columns=["path", "line", "sample"])
    measure = image_measures.iloc[0]

    p = Point(lon, lat)
    new_measures = []

    # list of matching results in the format:
    # [measure_index, x_offset, y_offset, offset_magnitude]
    match_results = []
    # lazily iterate for now
    for k,m in image_measures.iterrows():
        base_image = GeoDataset(m["path"])

        sx, sy = m["sample"], m["line"]

        for i,image in images.iterrows():
            # When grounding to THEMIS the df has a PATH to the QUAD
            dest_image = GeoDataset(image["path"])

            if os.path.basename(m['path']) == os.path.basename(image['path']):
                continue

            try:
                log.info(f'prop point: base_image: {base_image}')
                log.info(f'prop point: dest_image: {dest_image}')
                log.info(f'prop point: (sx, sy): ({sx}, {sy})')
                x,y, dist, metrics, corrmap = geom_match_simple(base_image, dest_image, sx, sy, 16, 16, \
                        match_func = match_func, \
                        match_kwargs=match_kwargs, \
                        verbose=verbose)
            except Exception as e:
                raise Exception(e)
                match_results.append(e)
                continue

            match_results.append([k, x, y,
                                 metrics, dist, corrmap, m["path"], image["path"],
                                 image['id'], image['serial']])

    # get best offsets
    match_results = np.asarray([res for res in match_results if isinstance(res, list) and all(r is not None for r in res)])
    if match_results.shape[0] == 0:
        # no matches
        return new_measures

    # column index 3 is the metric returned by the geom matcher
    best_results = np.asarray([match for match in match_results if cost(match_results[:,3], match[3])])
    if best_results.shape[0] == 0:
        # no matches satisfying cost
        return new_measures

    if verbose:
        log.info("match_results final length: ", len(match_results))
        log.info("best_results length: ", len(best_results))
        log.info("Full results: ", best_results)
        log.info("Winning CORRs: ", best_results[:,3], "Themis Pixel shifts: ", best_results[:,4])
        log.info("Themis Images: ", best_results[:,6], "CTX images:", best_results[:,7])
        log.info("Themis Sample: ", sx, "CTX Samples: ", best_results[:,1])
        log.info("Themis Line: ", sy, "CTX Lines: ", best_results[:,2])
        log.info('\n')

    # if the single best results metric (returned by geom_matcher) is None
    if len(best_results[:,3])==1 and best_results[:,3][0] is None:
        return new_measures

    height = dem.get_height(lat, lon)

    semi_major = config['spatial']['semimajor_rad']
    semi_minor = config['spatial']['semiminor_rad']
    # The CSM conversion makes the LLA/ECEF conversion explicit
    # reprojection takes ographic lat
    lon_og, lat_og = oc2og(lon, lat, semi_major, semi_minor)
    x, y, z = reproject([lon_og, lat_og, height],
                         semi_major, semi_minor,
                         'latlon', 'geocent')

    for row in best_results:
        sample = row[1]
        line = row[2]

        new_measures.append({
            'pointid' : pointid,
            'imageid' : row[8],
            'serial' : row[9],
            'path': row[7],
            'line' : line,
            'sample' : sample,
            'template_metric' : row[3],
            'template_shift' : row[4],
            'point' : p,
            'point_ecef' : Point(x, y, z)
            })

    return new_measures

def propagate_control_network(Session,
        config,
        dem,
        base_cnet,
        size_x=40,
        size_y=40,
        match_func="classic",
        match_kwargs={'image_size': (39,39), 'template_size': (21,21)},
        verbose=False,
        cost=lambda x,y: y == np.max(x)):
    """
    Loops over a base control network's measure information (line, sample, image path) and uses image matching
    algorithms (autocnet.matcher.subpixel.geom_match) to find the corresponding line(s)/sample(s) in database images.

    Parameters
    ----------
    Session : sqlalchemy.sessionmaker
              session maker associated with the database containing the images you want to propagate to

    config : dict
             configuation file associated with database containing the images you want to propagate to
             In the form: {'username':'somename', 'password':'somepassword',
             'host':'somehost', 'pgbouncer_port':6543, 'name':'somename'}

    dem : surface
          surface model of target body

    base_cnet : pd.DataFrame
                Dataframe representing the points you want to propagate. Must contain 'line', 'sample' location of
                the measure and the 'path' to the corresponding image

    verbose : boolean
              Increase the level of print outs/plots recieved during propagation

    cost : anonymous function
           determines to which image(s) the point should be propagated. x corresponds to a list
           of all match correlation metrics, while y corresponds to each indiviudal element
           of the x array.
           Example:
           cost = lambda x,y: y == np.max(x) will get you one result corresponding to the image that
           has the maximum correlation with the source image
           cost = lambda x,y: y > 0.6 will propegate the point to all images whose correlation
           result is greater than 0.6


    Returns
    -------
    ground : pd.DataFrame
             Dataframe containing pointid, imageid, image serial number, line, sample, and ground location (both latlon
             and cartesian) of successfully propagated points

    """
    log.warning('This function is not well tested. No tests currently exist \
    in the test suite for this version of the function.')

    match_func = check_match_func(match_func)

    groups = base_cnet.groupby('pointid').groups

    # append CNET info into structured Python list
    constrained_net = []

    # easily parallelized on the cpoint level, dummy serial for now
    for cpoint, indices in groups.items():
        measures = base_cnet.loc[indices]
        measure = measures.iloc[0]

        p = measure.point

        # get image in the destination that overlap
        lon, lat = measures["point"].iloc[0].xy
        gp_measures = propagate_point(Session,
                                      config,
                                      dem,
                                      lon[0],
                                      lat[0],
                                      cpoint,
                                      measures["path"],
                                      measures["line"],
                                      measures["sample"],
                                      size_x,
                                      size_y,
                                      match_func,
                                      match_kwargs,
                                      verbose=verbose,
                                      cost=cost)

        # do not append if propagate_point is unable to find result
        if len(gp_measures) == 0:
            continue
        constrained_net.extend(gp_measures)

    ground = gpd.GeoDataFrame.from_dict(constrained_net).set_geometry('point')
    groundpoints = ground.groupby('pointid').groups

    points = []

    # conditionally upload a new point to DB or updated existing point with new measures
    lat_srid = config['spatial']['latitudinal_srid']
    session = Session()
    for gp_point, indices in groundpoints.items():
        point = ground.loc[indices].iloc[0]

        # check DB for if point already exists there
        lon = point.point.x
        lat = point.point.y

        spatial_point = functions.ST_Point(lon, lat)
        spatial_setSRID = functions.ST_SetSRID(spatial_point, lat_srid)
        spatial_buffer = functions.ST_Buffer(spatial_setSRID, 10e-10)
        spatial_intersects = functions.ST_Intersects(Points.geom, spatial_buffer)

        res = session.query(Points).filter(spatial_intersects).all()

        if len(res) > 1:
           log.warning(f"There is more than one point at lon: {lon}, lat: {lat}")

        elif len(res) == 1:
            # update existing point with new measures
            for i in indices:
                row = ground.loc[i]
                pid = res[0].id
                meas = session.query(Measures.serial).filter(Measures.pointid == pid).all()
                serialnumbers = [m[0] for m in meas]

                if row['serial'] in serialnumbers:
                    continue

                points.append(Measures(pointid = pid,
                                       line = float(row['line']),
                                       sample = float(row['sample']),
                                       aprioriline = float(row['line']),
                                       apriorisample = float(row['sample']),
                                       imageid = int(row['imageid']),
                                       template_metric = float(row['template_metric']),
                                       template_shift = float(row['template_shift']),
                                       serial = row['serial'],
                                       measuretype = 3))
        else:
            # upload new point
            p = Points()
            p.pointtype = 3
            p.apriori = point['point_ecef']
            p.adjusted = point['point_ecef']
            for i in indices:
                row = ground.loc[i]
                p.measures.append(Measures(line = float(row['line']),
                                           sample = float(row['sample']),
                                           aprioriline = float(row['line']),
                                           apriorisample = float(row['sample']),
                                           imageid = int(row['imageid']),
                                           serial = row['serial'],
                                           template_metric = float(row['template_metric']),
                                           template_shift = float(row['template_shift']),
                                           measuretype = 3))
            points.append(p)

    session.add_all(points)
    session.commit()
    session.close()

    return ground


