#! /usr/bin/env python

import argparse

import numpy as np
from sklearn import mixture
import matplotlib.pyplot as plt

import storage, ml

parser = argparse.ArgumentParser(description='The program learns and displays Gaussians fit to beginnings and ends of object trajectories (based on Mohamed Gomaa Mohamed 2015 PhD). TODO: save the data')
parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file', required = True)
parser.add_argument('-t', dest = 'trajectoryType', help = 'type of trajectories to display', choices = ['feature', 'object'], default = 'object')
parser.add_argument('-norigins', dest = 'nOriginClusters', help = 'number of clusters for trajectory origins', required = True, type = int)
parser.add_argument('-ndestinations', dest = 'nDestinationClusters', help = 'number of clusters for trajectory destinations (=norigins if not provided)', type = int)
parser.add_argument('--covariance-type', dest = 'covarianceType', help = 'type of covariance of Gaussian model', default = "full")
parser.add_argument('-w', dest = 'worldImageFilename', help = 'filename of the world image')
parser.add_argument('-u', dest = 'unitsPerPixel', help = 'number of units of distance per pixel', type = float, default = 1.)
parser.add_argument('--display', dest = 'display', help = 'display points of interests', action = 'store_true') # default is manhattan distance

args = parser.parse_args()

objects = storage.loadTrajectoriesFromSqlite(args.databaseFilename, args.trajectoryType)

beginnings = []
ends = []
for o in objects:
    beginnings.append(o.getPositionAt(0).aslist())
    ends.append(o.getPositionAt(int(o.length())-1).aslist())

beginnings = np.array(beginnings)
ends = np.array(ends)

nDestinationClusters = args.nDestinationClusters
if args.nDestinationClusters is None:
    nDestinationClusters = args.nOriginClusters

gmmId=0
for nClusters, points, gmmType in zip([args.nOriginClusters, nDestinationClusters],
                                   [beginnings, ends],
                                   ['beginning', 'end']):
    # estimation
    gmm = mixture.GMM(n_components=nClusters, covariance_type = args.covarianceType)
    model=gmm.fit(beginnings)
    if not model.converged_:
        print('Warning: model for '+gmmType+' points did not converge')
    # plot
    if args.display:
        fig = plt.figure()
        if args.worldImageFilename is not None and args.unitsPerPixel is not None:
            img = plt.imread(args.worldImageFilename)
            plt.imshow(img)
        labels = ml.plotGMMClusters(model, points, fig, nUnitsPerPixel = args.unitsPerPixel)
        plt.axis('image')
        plt.title(gmmType)
        print(gmmType+' Clusters:\n{}'.format(ml.computeClusterSizes(labels, range(model.n_components))))
    # save
    storage.savePOIs(args.databaseFilename, model, gmmType, gmmId)
    gmmId += 1

if args.display:
    plt.show()

# fig = plt.figure()
# if args.worldImageFilename is not None and args.pixelsPerUnit is not None:
#     img = plt.imread(args.worldImageFilename)
#     plt.imshow(img)
# ml.plotGMMClusters(, , fig, nPixelsPerUnit = args.pixelsPerUnit)
# plt.axis('equal')
# plt.title()
# print('Destination Clusters:\n{}'.format(ml.computeClusterSizes(endModel.predict(ends), range(args.nClusters))))