import mandelbrot_1
import mandelbrot_2
import mandelbrot_3

import matplotlib.pyplot as plt
import numpy as np
import sys


"""Returns an array with the escape times for the provided rectangle,resotuion and max_escape_time
If filename provided, a plot will be made, shown and saved"""
def compute_mandelbrot(xmin,xmax,ymin,ymax,Nx,Ny,max_escape_time=1000,
                        plot_filename=None,cur_cmap=plt.cm.prism,totime=False,
                        mandel_func=mandelbrot_3.compute_mandelbrot):

    """
    Checks if rectangle is outside mandelbrot set before iterations.
    Checks if the circle that is made from abs(z) > 2, collides with the rectangle made from xmin,xmax,ymin,ymax

    Code based on : http://stackoverflow.com/a/21096179

    """
    def is_not_outside(xmin,xmax,ymin,ymax):

        rect_x = xmin   #corner coordinate x
        rect_y = ymin   #corner coordinate y
        rect_w = abs(xmax-xmin) #width
        rect_h = abs(ymax-ymin) #height

        #abs > 2 is a circle at 0,0 with radius 2
        circle_x = 0
        circle_y = 0
        circle_radius = 2

        distX = abs(circle_x - rect_x - rect_w/2)
        distY = abs(circle_y - rect_y - rect_h/2)

        dist_half_w = (rect_w/2 + circle_radius)
        dist_half_h = (rect_h/2 + circle_radius)

        if (distX > dist_half_w) or (distY > dist_half_h) :
            return False

        if (distX <= (rect_w/2)) or (distY <= (rect_h/2)):
            return True

        dist_center_x = distX - rect_w/2
        dist_center_y = distY - rect_h/2

        if dist_center_x**2 + dist_center_y**2 <= (circle_radius**2):
            return True
        else:
            return False


    if is_not_outside(xmin,xmax,ymin,ymax):
        return mandel_func(xmin,xmax,ymin,ymax,Nx,Ny,max_escape_time,
                        plot_filename,cur_cmap,totime)

    else:
        raise ValueError("Rectangle is entirely outstide the mandelbrot set. Invalid.")




script_name = __file__
usage_string = 'USAGE\t : ' + script_name + ' [rect=xmin,xmax,ymin,ymax] [res=Nx,Ny] [esc=INTEGER] [ver=1|2|3] [col=1|2|3] [out=FILENAME] [totime]'

"""Prints usage output and exits execution"""
def usage(error):
    print(error)
    print(usage_string)
    print(script_name + ': try \'' + script_name + ' --help\' for more information')
    sys.exit(0)

"""Prints help information when called with --help. Exits after"""
def help():
        print(usage_string)
        print(' Explanation:')
        # print()
        print('  rect=xmin,xmax,ymin,ymax  : Sets rectangle to draw to [xmin,xmax] x [ymin,ymax] \n\tDefault : rect=-2,1,-1,-1')
        print('  res=Nx,Ny                 : Sets resolution to Nx x Ny (both positive integers) \n\tDefault : 200 x 200')
        print('  esc=INTEGER               : Sets max escape time provided integer \n\tDefault : 1000')
        print('  ver=1|2|3                 : Choses impletation : 1 -> 4.1 , 2 -> 4.2 , 3 -> 4.3 \n\tDefault: 2 (4.2)')
        print('  col=1|2|3                 : Choses color scale : 1 -> flag , 2 -> prism , 3 -> jet \n\tDefault: flag')
        print('  out=FILENAME              : Filename to store the drawn image \n\tDefault: None')
        print('  totime                    : Turns on printing of timing \n\tDefault: Set to False')
        print('  test                      : Test-runs all the 3 implementations with timing \n\tDefault: Set to False')
        sys.exit(0)


def run_from_cmdline():


    #default values
    xmin = -2
    xmax = 1
    ymin = -1
    ymax = 1
    Nx = 400
    Ny = 400
    max_escape_time = 500
    ver = 3
    col = 1
    plot_filename = None

    timeflag = False
    testflag = False

    args = sys.argv[1:] #skip file name
    for arg in args:
        arg_split = arg.split('=')
        try:
            if arg == '--help':
                help()

            elif arg == 'totime':   #timing
                timeflag = True
                continue

            elif arg == 'test':     #added functionality
                testflag = True
                continue

            elif (len(arg_split) < 2):
                usage("Invalid input: " + arg)

            elif arg_split[0] == "rect":
                numbers_str = arg_split[1].split(',')
                xmin,xmax,ymin,ymax = [float(x) for x in numbers_str]
                if(xmin > xmax):
                    usage("xmin >= xmax, invalid")
                if(ymin > ymax):
                    usage("ymin >= ymax, invalid")


            elif arg_split[0] == "res":
                numbers_str = arg_split[1].split(',')
                intvalNx,intvalNy = [int(x) for x in numbers_str]
                if intvalNx > 1 or intvalNy > 1:
                    Nx = intvalNx
                    Ny = intvalNy
                else:
                    usage("resolution invalid : {},{} . Should be positive integers ".format(intvalNx,intvalNy))


            elif arg_split[0] == "esc":
                intval = int(arg_split[1])
                if intval >= 1 :
                    max_escape_time = intval
                else:
                    usage("resolution invalid : {} . Should be positive integer ".format(intvalNx))

            elif arg_split[0] == 'ver':
                intval = int(arg_split[1])
                if  (1 <= intval <= 3):
                    ver = int(arg_split[1])
                else:
                    usage("Implementation input invalid : {} . Should be 1,2 or 3 ".format(intval))

            elif arg_split[0] == "out":
                plot_filename = arg_split[1]

            elif arg_split[0] == 'col':
                intval = int(arg_split[1])
                if  (1 <= intval <= 3):
                    col = intval
                else:
                    usage("Color scale invalid : {} . Should be 1,2 or 3 ".format(intval))

            else:
                usage('INVALID OPTION: ' + arg)

        except ValueError:
            usage(arg_split[1] + " invalid input type or form for " + arg_split[0])
        except IndexError:
            usage(arg_split[1] + " : is not enough numbers in input for " + arg_split[0])




    #set the function from right implementation
    mandel_func =  mandelbrot_3.compute_mandelbrot      #default
    if ver == 1:
        mandel_func = mandelbrot_1.compute_mandelbrot
    elif ver == 2:
        mandel_func = mandelbrot_2.compute_mandelbrot

    #set the right color map
    color_scale = plt.cm.flag               #default
    color_scale.set_bad('black',1.)
    if col == 2:
        color_scale = plt.cm.prism
        color_scale.set_bad('black',1.)
    elif col == 3:
        color_scale = plt.cm.jet



    if(not testflag):   #normal
        compute_mandelbrot(xmin,xmax,ymin,ymax,Nx,Ny,max_escape_time,plot_filename,cur_cmap=color_scale,
            totime=timeflag,mandel_func=mandel_func)
    else:
        #Runs timing for all 3 implementations and compares return-arrays
        print("Testing all implementations")
        print("args: ",Nx,Ny,xmin ,xmax, ymin ,ymax ,max_escape_time,ver,col)
        a1 = mandelbrot_1.compute_mandelbrot(xmin,xmax,ymin,ymax,Nx,Ny,max_escape_time,None,cur_cmap=None,totime=True)
        a2 = mandelbrot_2.compute_mandelbrot(xmin,xmax,ymin,ymax,Nx,Ny,max_escape_time,None,cur_cmap=None,totime=True)
        a3 = mandelbrot_3.compute_mandelbrot(xmin,xmax,ymin,ymax,Nx,Ny,max_escape_time,None,cur_cmap=None,totime=True)
        print("4.1 == 4.2 ->",np.array_equal(a1, a2))
        print("4.1 == 4.3 ->",np.array_equal(a1, a3))
        print("4.2 == 4.3 ->",np.array_equal(a2, a3))





if __name__ == "__main__":
    try:
        run_from_cmdline()
    except ValueError as e: #catch ValueError if rectangle is enteriely outside from the start
        print(e)
        sys.exit(1)