import matplotlib.pyplot as plt
import numpy as np
import cv2 as cv
import os
import pyskyline

import pyskyline.localization

FILE = 'terrain'
SIZE = 256
X    = SIZE//2
Y    = SIZE//2+80

if __name__ == '__main__':
    if os.path.isfile(f'{FILE}_{SIZE}.npz'):
        print("Save detected, loading ...")
        terrain = pyskyline.Terrain.load(f'{FILE}_{SIZE}')
    else:
        print("No save detected, generating new terrain ...")
        terrain = pyskyline.Terrain.generate(size=SIZE)
        terrain.compute_all_skylines()
        terrain.save(f'{FILE}_{SIZE}')
        print(f"Done, terrain saved as {FILE}_{SIZE}.npz")

    # Output elevation map
    cv.imwrite('elevation_map.jpg', terrain.dem.elevation)

    # Output color map
    cv.imwrite('color_map.jpg', terrain.dem.color)

    # Ouput field of view visualization
    cv.imwrite('fov.jpg', terrain.compute_fov(X, Y))

    # Ouput skyline
    skyline = terrain.compute_skyline(X, Y)
    plt.figure()
    plt.grid(False)
    plt.plot(np.arange(-180,180,360/256),skyline)
    plt.title('Skyline')
    plt.savefig('skyline.svg')

    # Generate heatmap
    skyline  = np.roll(skyline, 8)           # Add heading error
    skyline += np.random.normal(0,0.5, 256)  # Add random noise
    mosse_correlation = pyskyline.MosseCorrelation(terrain, skyline)
    scoremap = mosse_correlation.generate_scoremap()

    plt.figure()
    plt.grid(False)
    fig = plt.imshow(scoremap)
    plt.title('Score map')
    plt.colorbar(fig)
    plt.savefig('scoremap.svg')