# copyright ############################### #
# This file is part of the Xtrack Package.  #
# Copyright (c) CERN, 2021.                 #
# ######################################### #

import numpy as np
import matplotlib.patches
import matplotlib.collections
import matplotlib.pyplot as plt
#plt.style.use('kostas')

def initial_xy_polar(r_min, r_max, r_N, theta_min, theta_max, theta_N):
    return np.array([[(r*np.cos(theta),r*np.sin(theta)) for r in np.linspace(r_min,r_max,r_N)] for theta in np.linspace(theta_min,theta_max,theta_N)])

def initial_xy_cartesian(x_min, x_max, x_N, y_min, y_max, y_N):
    return np.array([[(x,y) for x in np.linspace(x_min,x_max,x_N)] for y in np.linspace(y_min,y_max,y_N)])


def draw_footprint(A, axis_object=None, figure_object=None, axis=0, linewidth=4):
    '''
    Input A should be a 3-D numpy array with shape (Nx,Ny,2)
    representing a 2-D array of (x,y) points. This function
    will draw lines between adjacent points in the 2-D array.
    '''
    if len(A.shape) != 3:
        print('ERROR: Invalid input matrix')
        return None
    if A.shape[2] != 2:
        print('ERROR: Points are not defined in 2D space')
        return None

    sx = A.shape[0]-1
    sy = A.shape[1]-1

    p1 = A[:-1,:-1,:].reshape(sx*sy,2)[:,:]
    p2 = A[1:,:-1,:].reshape(sx*sy,2)[:]
    p3 = A[1:,1:,:].reshape(sx*sy,2)[:]
    p4 = A[:-1,1:,:].reshape(sx*sy,2)[:]

    #Stack endpoints to form polygons
    Polygons = np.stack((p1,p2,p3,p4))
    #transpose polygons
    Polygons = np.transpose(Polygons,(1,0,2))
    patches = list(map(matplotlib.patches.Polygon,Polygons))

    #assign colors
    patch_colors = [(0,0,0) for a in Polygons]
    patch_colors[(sx-1)*sy:] = [(0,1,0)]*sy
    patch_colors[(sy-1)::sy] = [(0,0,1)]*sx

    p_collection = matplotlib.collections.PatchCollection(patches,facecolors=[],linewidth=linewidth,edgecolor=patch_colors)

    if axis_object is None:
        if figure_object:
            fig = figure_object
        else:
            fig = plt.figure()
        if len(fig.axes) == 0:
            plt.subplot(1,1,1)
        if axis >= len(fig.axes) or axis < 0:
            i = 0
        else:
            i = axis
        ax = fig.axes[i]
    else:
        ax = axis_object
        fig = None

    ax.add_collection(p_collection)

    return fig

def example():
    coords_matrix = initial_xy_polar(0.01,5,11,0.,np.pi/2.,10)
    fig = draw_footprint(coords_matrix)
    fig.axes[0].set_xlim(-0.05,5.1)
    fig.axes[0].set_ylim(-0.05,5.1)
