"""Module for the simulation of the methods."""

from __future__ import print_function, division
import sys
import logging
from math import atan2, pi, sqrt
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from os.path import isfile

import numpy as np
from numpy.random import random
import matplotlib.pyplot as plt
import shapely.geometry as geom
from scipy.optimize import curve_fit

# General
logger = None
LOG_MSG_FMT = "[%(asctime)s][%(levelname)-8s]\
[%(filename)s, %(lineno)d][%(name)s]\t%(message)s"
LOG_DT_FMT = "\033[1m%m-%d %H:%M:%S\033[0m"

CMDS = ['position', 'size']
USAGE_STR = "<command> [options]"
EXAMPLES_CMDS = """
\033[1mEXAMPLES\033[0m
$ python method_sims.py position [options]
$ python method_sims.py size [options]
"""

# Beam position
CENTER = (1, 0)
RADIUS = 1
NUM = 1000
DEMO_TRACKS = [[(1, 0), (1-sqrt(2)/2, sqrt(2)/2)],
               [(1, 0), (1+sqrt(2)/2, sqrt(2)/2)],
               [(1, 0), (1-sqrt(2)/2, -sqrt(2)/2)],
               [(1, 0), (1+sqrt(2)/2, -sqrt(2)/2)],
               [(1, 0), (0, 0)], [(1, 0), (2, 0)],
               [(1, 0), (1, 1)], [(1, 0), (1, -1)]
               ]


def set_logging(verbosity, msg_fmt=LOG_MSG_FMT, dt_fmt=LOG_DT_FMT):
    """Sets the logging output style.
    """
    logging.addLevelName(logging.DEBUG, "\033[1;34m%-8s\033[1;0m"
                         % logging.getLevelName(logging.DEBUG))
    logging.addLevelName(logging.INFO, "\033[1;37m%-8s\033[1;0m"
                         % logging.getLevelName(logging.INFO))
    logging.addLevelName(logging.WARNING, "\033[1;33m%-8s\033[1;0m"
                         % logging.getLevelName(logging.WARNING))
    logging.addLevelName(logging.ERROR, "\033[1;31m%-8s\033[1;0m"
                         % logging.getLevelName(logging.ERROR))
    logging.addLevelName(logging.CRITICAL, "\033[1;41m%-8s\033[1;0m"
                         % logging.getLevelName(logging.CRITICAL))

    logging.basicConfig(level=verbosity, format=LOG_MSG_FMT,
                        datefmt=LOG_DT_FMT)


def counts_to_verbosity(counts):
    """Gets the 'v' count from the cli and translates it to some logging level
    """
    if counts is None:
        return logging.WARNING
    elif counts == 1:
        return logging.INFO
    else:
        return logging.DEBUG


def doca(phi, x0, y0):
    """DOCA formula."""
    return x0*np.sin(phi) - y0*np.cos(phi)


def doca2(phi, x0, y0):
    """DOCA formula."""
    return y0*np.cos(phi) - x0*np.sin(phi)


class SimGameBGV(object):
    """Play simulating the methods for beam position and size measurement."""

    def __init__(self):
        """Init mostly for ArgParser and command definition."""
        parser = ArgumentParser(description=self.__doc__,
                                formatter_class=RawDescriptionHelpFormatter,
                                prog='method_sims.py',
                                epilog=EXAMPLES_CMDS)
        parser.add_argument('command', help='Subcommand to run %s' % CMDS)

        args = parser.parse_args(sys.argv[1:2])
        if not hasattr(self, args.command):
            logger.error('ERROR: Unrecognized command')
            parser.print_help()
            exit(1)

        getattr(self, args.command)()

    def set_position_parser(self):
        parser = ArgumentParser(description='Simulate position calculation',
                                prog='method_sims.py position')

        parser.add_argument('-v', '--verbosity', action='count',
                            help='Increase output verbosity')
        parser.add_argument('--pres', action='store_true', default=False,
                            help="Predefined lims etc for presentation")

        parser.add_argument('--num', '-n', action='store', default=NUM,
                            type=int,
                            help='Number of points/tracks to create')
        parser.add_argument('--def', action='store_true', default=False,
                            help="If predefined tracks are used - Hardcoded.")
        parser.add_argument('--defindex', action='store', type=str, default="",
                            help="a list with the indexes to be used for demo\
                            like 0,1,3,5")
        parser.add_argument('--both', action='store_true', default=False,
                            help="Both random and predefined tracks.")

        parser.add_argument('--pos', '-p', action='store', default=(1, 0),
                            type=tuple, help="Tuple of the center position")
        parser.add_argument('--sigma', '-s', action='store', default=0.0,
                            type=float,
                            help='Width of distribution for track creation\
                            for the center point of the tracks')
        parser.add_argument('--radius', '-r', action='store', default=RADIUS,
                            type=float,
                            help='Radius of the created random points to make \
                            the projected track')
        parser.add_argument('--offset', '-f', action='store', default=0,
                            type=float,
                            help='Offset of the points to create the tracks')

        parser.add_argument('--noposfit', action='store_false', default=True,
                            help="Fit the data with DOCA ?")
        parser.add_argument('--save', action='store_true', default=False,
                            help="Save the figures.")

        return parser

    def set_size_parser(self):
        parser = ArgumentParser(description='Simulate size calculation [IPC]',
                                prog='method_sims.py size')

        parser.add_argument('-v', '--verbosity', action='count',
                            help='Increase output verbosity')

        parser.add_argument('--num', '-n', action='store', default=NUM,
                            type=int,
                            help='Number of points/tracks to create')

        parser.add_argument('--pos', '-p', action='store', default=(1, 0),
                            type=tuple, help="Tuple of the center position")
        parser.add_argument('--sigma', '-s', action='store', default=0.0,
                            type=float,
                            help='Width of distribution for track creation\
                            for the center point of the tracks')
        parser.add_argument('--radius', '-r', action='store', default=RADIUS,
                            type=float,
                            help='Radius of the created random points to make \
                            the projected track')
        parser.add_argument('--offset', '-f', action='store', default=0,
                            type=float,
                            help='Offset of the points to create the tracks')

        parser.add_argument('--numtracks', '-t', action='store', default=3,
                            type=int,
                            help='Input to randint function - define how many\
                            tracks per event to be constructed. First param\
                            always 2, here only the open upper limit.')

        return parser

    def position(self, **kwargs):
        parser = self.set_position_parser()

        #######################################################################
        # INIT
        #######################################################################
        args, other = parser.parse_known_args(sys.argv[2:])
        # args = vars(parser.parse_args(sys.argv[2:]))
        args = vars(args)
        others = dict(zip(other[::2], other[1::2]))

        set_logging(counts_to_verbosity(args['verbosity']))
        global logger
        logger = logging.getLogger(__name__)

        center = args['pos'][0], args['pos'][1]
        ctr_axes = geom.Point(0.0, 0.0)

        tracks = []
        phis = []
        ips = []
        fit = []

        logger.info('Running \'POSITION\' command')
        #######################################################################
        # BUILDING POINTS BASED ON ARGS
        #######################################################################
        cntr_x, cntr_y, px, py = [], [], [], []
        if args['def'] or args['both']:
            logger.debug('Using pre-defined tracks.')

            demo_tracks = []
            if args['defindex']:
                defind = args['defindex'].split(',')
                for item in defind:
                    try:
                        ind = int(item)
                        if ind >= len(DEMO_TRACKS):
                            logger.warning('Index %s larger than demo length')
                        else:
                            demo_tracks.append(DEMO_TRACKS[ind])
                    except ValueError:
                        logger.warning('Wrong index input: %s' % item)
            else:
                demo_tracks = DEMO_TRACKS

            logger.debug('Demo tracks to be used: %s' % demo_tracks)
            cntr_x = [item[0][0] for item in demo_tracks]
            cntr_y = [item[0][1] for item in demo_tracks]
            px = [item[1][0] for item in demo_tracks]
            py = [item[1][1] for item in demo_tracks]

            n = range(len(cntr_x))

        if args['both'] or not args['def']:
            logger.debug('Using random tracks.')
            logger.debug('Center point: %s' % str(center))
            n = range(args['num'])
            # Create center point - beam position
            if args['sigma']:
                w = args['sigma']
                logger.debug('With noise of sigma= %s' % w)
                cntr_x_tmp = np.random.normal(loc=center[0], scale=w,
                                              size=len(n))
                cntr_y_tmp = np.random.normal(loc=center[1], scale=w,
                                              size=len(n))
            else:
                logger.debug('As a singular point')
                cntr_x_tmp = np.repeat(center[0], len(n))
                cntr_y_tmp = np.repeat(center[1], len(n))

            # Create secondary point for track to be defined
            px_tmp = (args['radius'] * random(size=len(n)) + args['offset']) *\
                     np.cos(2 * pi*random(size=len(n))) + cntr_x_tmp
            py_tmp = (args['radius'] * random(size=len(n)) + args['offset']) *\
                     np.cos(2 * pi*random(size=len(n))) + cntr_y_tmp

            cntr_x += list(cntr_x_tmp)
            cntr_y += list(cntr_y_tmp)
            px += list(px_tmp)
            py += list(py_tmp)

        #######################################################################
        # LOOP
        #######################################################################
        logger.debug('Loops number: %d' % len(n))
        for i in n:
            # Angles - calc the phi of the line
            dy, dx = py[i]-cntr_y[i], px[i]-cntr_x[i]
            ext_coord = []
            if dy == 0:
                b = cntr_y[i]
                ext_coord = [(0, b), (px[i], b)]
            elif dx == 0:
                b = cntr_x[i]
                ext_coord = [(b, 0), (b, py[i])]
            else:
                b = cntr_y[i]-(dy/dx)*cntr_x[i]
                ext_coord = [(0, b), (-b/(dy/dx), 0)]

            phi = atan2(dy, dx)
            phis.append(phi)

            # Finally store info and create the track
            track = geom.LineString(((cntr_x[i], cntr_y[i]), (px[i], py[i])))
            tracks.append(track)

            # Distance (IP)
            # extend track so that it touches the axes (needed? --> yes)
            line = geom.LineString(ext_coord)
            if phi:
                # ip = ctr_axes.distance(line)*(phi/abs(phi))   # same ofc
                ip = line.distance(ctr_axes)*(phi/abs(phi))
            else:
                ip = 0
            ips.append(ip)
            # TODO - DOCA but need to find POCA first to be the point
            # ip.append(point.y*np.cos(ang)-point.x*np.sin(ang))

        # FIT -- after the loop
        if args['noposfit'] and len(ips) > 1:
            logger.debug('Fitting the DOCA-Phi...')
            phis = np.array(phis)
            ips = np.array(ips)
            ret = curve_fit(doca, phis, ips)
            res = ret[0]
            logger.info('Fit params:')
            logger.info('\t(x,y)=(%f, %f)' % (np.round(res[0], 2),
                                              np.round(res[1], 2)))
            logger.info('\tcovM= %s' % ret[1])
            xi = np.linspace(-pi, pi)
            yi = doca(xi, res[0], res[1])
            fit = {'data': [xi, yi], 'coeffs': [res[0], res[1]]}

        kw = {**args, **others}
        beam_pos_plotting(tracks, phis, ips, fit, **kw)

    def size(self, **kwargs):
        #######################################################################
        # INIT
        #######################################################################
        parser = self.set_size_parser()
        args, other = parser.parse_known_args(sys.argv[2:])
        # args = vars(parser.parse_args(sys.argv[2:]))
        args = vars(args)

        set_logging(counts_to_verbosity(args['verbosity']))
        global logger
        logger = logging.getLogger(__name__)

        center = args['pos'][0], args['pos'][1]
        ctr_axes = geom.Point(0.0, 0.0)

        tracks = []
        phis = []
        ips = []
        fit = []

        logger.info('Running \'SIZE\' command')
        logger.debug('Using random tracks.')
        logger.debug('Center point: %s' % str(center))
        n = range(args['num'])
        # Create center point - beam position
        if args['sigma']:
            w = args['sigma']
            logger.debug('With noise of sigma= %s' % w)
            cntr_x = np.random.normal(loc=center[0], scale=w, size=len(n))
            cntr_y = np.random.normal(loc=center[1], scale=w, size=len(n))
        else:
            logger.debug('As a singular point')
            cntr_x = np.repeat(center[0], len(n))
            cntr_y = np.repeat(center[1], len(n))

        # Create secondary point for track to be defined
        px = (args['radius'] * random(size=len(n)) + args['offset']) *\
            np.cos(2 * pi*random(size=len(n))) + cntr_x
        py = (args['radius'] * random(size=len(n)) + args['offset']) *\
            np.cos(2 * pi*random(size=len(n))) + cntr_y
        # TODO - create more tracks of same vertex / correct angles etc


def beam_pos_plotting(tracks, phis, ips, fit, **kwargs):
    """Plot the results for the beam position method."""
    # onlypoints = int(kwargs.get('onlypoints', 0))
    m1 = kwargs.get('pmarker', '.')
    m2 = kwargs.get('cmarker', '.')

    # Plotting of tracks or points (the random ones used for track creation)
    logger.debug('Plotting Tracks...')
    f1_title = 'Tracks'
    f1 = plt.figure(f1_title, figsize=(12, 8))
    ax1 = f1.add_subplot(111)
    """
    if onlypoints:
        logger.debug('Plotting points...')
        ox = [t.coords[:][1][0] for t in tracks]
        oy = [t.coords[:][1][1] for t in tracks]
        ax1.plot(ox, oy, ls='None', marker=m1, mfc='k')
    """
    for t in tracks:
        ax1.plot(*t.xy, lw=2)

    if kwargs['sigma']:
        cx = [t.coords[:][0][0] for t in tracks]
        cy = [t.coords[:][0][1] for t in tracks]
        ax1.plot(cx, cy, ls='None', marker=m1, mfc='b')
    else:
        x, y = tracks[0].coords[:][0][0], tracks[0].coords[:][0][1]
        ax1.plot([x], [y], ls='None', marker='D', mfc='r', ms=15)
    ax1.tick_params(labelsize=16)
    ax1.set_xlabel('$x [mm]$', fontsize='30')
    ax1.set_ylabel('$y [mm]$', fontsize='30')
    ax1.grid(True, lw=2)
    ax1.set_aspect('equal')

    ctr_x, ctr_y = kwargs['pos'][0], kwargs['pos'][1]
    radius = kwargs['radius']
    ax1.set_xlim((ctr_x-radius, ctr_x+radius))
    ax1.set_ylim((-ctr_y-radius, ctr_y+radius))

    logger.debug('Plotting Phi histogram...')
    f2_title = 'Angles Hist'
    f2 = plt.figure(f2_title, figsize=(12, 8))
    ax2 = f2.add_subplot(111)
    ax2.hist(phis, bins=90, color='k', fill=True, alpha=.6, lw=2)
    ax2.tick_params(labelsize=16)
    ax2.set_xlabel('$\phi [rad]$', fontsize='30')
    ax2.set_xlim((-4, 4))
    # ax2.set_ylim((0, 2))

    logger.debug('Plotting DOCA-Phi...')
    f3_title = 'Position'
    f3 = plt.figure(f3_title, figsize=(12, 8))
    ax3 = f3.add_subplot(111)
    ax3.plot(phis, ips, ls='None', marker='x', mew=3, c='k', mfc='k', ms=10)
    if fit:
        xi, yi = fit['data']
        ax3.plot(xi, yi, ls='--', c='r', lw=4)
    ax3.tick_params(labelsize=16)
    ax3.set_xlabel('$\phi [rad]$', fontsize='30')
    ax3.set_ylabel('$d_{xy}$', fontsize='30')
    ax3.set_xlim((-4, 4))
    ax3.set_ylim((-1.2, 1.2))

    plt.show()

    if kwargs['save']:
        logger.debug('Saving figures...')
        ext = '.png'
        cnt = 0
        while isfile('%s_' % cnt + f1_title.lower() + ext):
            cnt += 1
        f1.savefig('%s_' % cnt + f1_title.lower() + ext)
        f2.savefig('%s_' % cnt + f2_title.lower() + ext)
        f3.savefig('%s_' % cnt + f3_title.lower() + ext)

    logger.debug('The end!')

if __name__ == '__main__':
    SimGameBGV()
