#! /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('-n', dest = 'nClusters', help = 'number of point clusters', required = True, 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 = 'pixelsPerUnit', help = 'number pixels per unit of distance', type = float, default = 1.)

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)

gmm = mixture.GMM(n_components=args.nClusters, covariance_type = args.covarianceType)
beginningModel=gmm.fit(beginnings)
gmm = mixture.GMM(n_components=args.nClusters, covariance_type = args.covarianceType)
endModel=gmm.fit(ends)

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(beginningModel, beginnings, fig, nPixelsPerUnit = args.pixelsPerUnit)
plt.axis('equal')
plt.title('Origins')
print('Origin Clusters:\n{}'.format(ml.computeClusterSizes(beginningModel.predict(beginnings), range(args.nClusters))))

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(endModel, ends, fig, nPixelsPerUnit = args.pixelsPerUnit)
plt.axis('equal')
plt.title('Destinations')
print('Destination Clusters:\n{}'.format(ml.computeClusterSizes(endModel.predict(ends), range(args.nClusters))))