#!/usr/bin/env python3

import math
import cv2 as cv
import pyskyline
import numpy as np
import argparse
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider

def show_result(terrain, skyline, coord, scoremap):
    fig = plt.figure()
    gs = fig.add_gridspec(2,3)
    fig.suptitle("Mosse correlation", fontsize=16)
    ax1 = fig.add_subplot(gs[0, 0])
    ax2 = fig.add_subplot(gs[0, 1])
    ax3 = fig.add_subplot(gs[0, 2])
    ax4 = fig.add_subplot(gs[1, :])

    im1 = ax1.imshow(scoremap, cmap='hot', interpolation='nearest')
    ax1.set_title('Correlation map')
    cbar1 = ax1.figure.colorbar(im1)
    cbar1.ax.set_ylabel("Correlation score", rotation=-90, va="bottom")

    rgb = cv.cvtColor(terrain.rgb(), cv.COLOR_BGR2RGB)
    rgb = cv.circle(rgb, coord, 2, (255,0,0), -1)
    ax2.imshow(rgb)
    ax2.set_title('Colored terrain')

    im2 = ax3.imshow((terrain.grayscale().astype(float) - 135) * (40.0/120.0))
    ax3.set_title('DEM')
    cbar = ax1.figure.colorbar(im2)
    cbar.ax.set_ylabel("Altitude", rotation=-90, va="bottom")

    x = np.linspace(-180,180,skyline.shape[0])
    ax4.plot(x,skyline,label="Agent vision")
    ax4.plot(x,terrain.get_skyline(coord[0], coord[1]), label="DEM data")
    ax4.legend()

    plt.show()

def eval_score(correlation_map, coord, scale, alpha=10):
    alpha = alpha / scale
    score = 0
    h,w   = correlation_map.shape
    x,y   = coord
    for j in range(h):
        for i in range(w):
            dist   = math.sqrt((x-i)**2+(y-j)**2)
            score += correlation_map[j,i] * 1 #(math.pow(0.5,math.pow(dist/alpha,2)-1)-1)
    return score

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument("--size",  help="Terrain size (default=128)",               type=int, default=128)
    parser.add_argument("--seed",  help="Terrain seed (If not set seed is random)", type=int)

    args = parser.parse_args()

    seed     = np.random.randint(0,1024) if args.seed is None else args.seed
    terrain  = pyskyline.Terrain.generate(size=args.size,seed=seed,scale=8.0,max_altitude=40.0)

    terrain.precompute_all_skylines()

    x, y            = terrain.get_random_water_coordinate()
    skyline         = terrain.get_skyline(x, y)
    angle           = np.random.randint(1,40)
    skyline         = np.roll(skyline, angle)
    correlation_map = pyskyline.MosseCorrelation.process(terrain, skyline)
    score           = eval_score(correlation_map,(x,y), terrain._scale,100)
    score           = score / (2 * terrain.water_tile_count()) + 0.5
    print(f"{score:0.2f} with angle of {angle}")
    show_result(terrain, skyline, (x,y), correlation_map)