from scipy import log, sqrt, linspace, logspace, meshgrid, zeros
from scipy.optimize import minimize


def isstable(model):
    """

    :param model: (F,G,H)
    :return: 2GH + H^2 - 2F, > 0 if stable; < 0 if unstable
    """
    return 2 * model[1] * model[2] + model[2] * model[2] - 2 * model[0]


def tflogmag((f, g, h), w):
    """
    Return the log of the magnitude of the linearized transfer function
        log[ |T(jw)| ]
    """
    # f,g,h = model.f, model.g, model.h
    return log(
        (g * g * w * w + f * f) / ((f - w * w) ** 2 + (g + h) ** 2 * w * w)) / 2


def htflogmag((f, g, h), w):
    """
    Return the log of the magnitude of the linearized headway transfer function
        log[ |1-T(jw)| ]
    """
    # f,g,h = model.f, model.g, model.h
    return log((h * h * w * w + w * w * w * w) / (
        (f - w * w) ** 2 + (g + h) ** 2 * w * w)) / 2


def kstable(human, robot, w, eta=2):
    """
    Return the maximum number of human cars 
        that can be string-stabilized by a single robot car
        for an oscillatory perturbation of frequency w
    """
    return -tflogmag(robot, w) / tflogmag(human, w)


def ksafe(human, robot, w, eta=2):
    """
    Return the maximum number of human cars
        that can be "safely" followed by a single robot car
        for an oscillatory perturbation of frequency w
        by a safety parameter eta
    :param human:
    :param robot:
    :param w:
    :param eta: safety margin
    :return:
    """
    return (log(eta) - htflogmag(robot, w)) / tflogmag(human, w)


TOL = 1e-6


def maxkstable(human, robot, eta=2):
    """
    Find the maximum number of humans cars that can be string-stabilized by a
    single robot car for any frequency (all w)
    :param human:
    :param robot:
    :param eta:
    :return:
    """
    f, g, h = human
    maxw = sqrt(2 * f - 2 * g * h - h * h)
    return minimize(lambda w: kstable(human, robot, w), 1,
                    bounds=[(TOL, maxw - TOL)])  # .fun[0]


def maxksafe(human, robot, eta=2):
    """
    Find the maximum number of humans cars that can be "safely" followed by a
    single robot car for any frequency (all w)
    :param human:
    :param robot:
    :param eta:
    :return:
    """
    f, g, h = human
    maxw = sqrt(2 * f - 2 * g * h - h * h)
    return minimize(lambda w: ksafe(human, robot, w, eta), 1,
                    bounds=[(TOL, maxw - TOL)])  # .fun[0]


def maxk(human, robot, eta):
    """
    Min of maxksafe and maxkstable
    :param human:
    :param robot:
    :param eta:
    :return:
    """
    return min(maxkstable(human, robot), maxksafe(human, robot, eta))


'''
def bestrobotstable(human):
    cons = ({'type': 'ineq', 'fun': lambda x: 2*x[1]*x[2] + x[2]*x[2] - 2*x[0] },)
    return minimize(lambda robot : -maxkstable(human, robot), 
            (0.1, 0, 0.1), 
            bounds=[(0, 0.2), (0, 0), (0, 0.2)],
            constraints=cons)

def bestrobotsafe(human, eta):
    cons = ({'type': 'ineq', 'fun': lambda x: 2*x[1]*x[2] + x[2]*x[2] - 2*x[0] },)
    return minimize(lambda robot : -maxksafe(human, robot, eta), 
            (0.1, 0, 0.1), 
            bounds=[(0, 0.2), (0, 0), (0, 0.2)],
            constraints=cons)
'''


def bestrobot(human, eta, fbound=(0, 0.2), gbound=(0, 0.2), hbound=(0, 0.2)):
    cons = ({'type': 'ineq', 'fun': lambda x: isstable(x)},)
    return minimize(lambda robot: -maxk(human, robot, eta).fun[0],
                    [0.001,0.001,0.001], constraints=cons,
                    bounds=[fbound, gbound, hbound])


if __name__ == "__main__":
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import pyplot as plt
    from cfm import IDM, CFM

    hmodel = IDM(a=0.3, b=3, t=1.5, s0=2, v0=30)
    hmodel.go(0.5)
    human = (hmodel.f, hmodel.g, hmodel.h)

    maxw = sqrt(2 * hmodel.f - 2 * hmodel.g * hmodel.h - hmodel.h * hmodel.h)
    # print "maxw = ", maxw
    resw = sqrt(-isstable(human))
    # print "resw = ", resw

    robot = (0.018, 0, 0.19)
    # print 2*robot[0] - 2*robot[1]*robot[2] - robot[2]*robot[2]

    eta = 2

    print bestrobot(human, eta)

    import IPython

    IPython.embed()

    '''
    ws = linspace(TOL, maxw-TOL, 101)
    kst = kstable(human, robot, ws)
    ksf = ksafe(human, robot, ws, eta)
    # print "kstable = ", kstable(human, robot, ws)
    plt.plot(ws, kst)
    plt.plot(ws, ksf)
    plt.axis([0, maxw, 0, 100])
    plt.show()

    frng = linspace(0, 0.1, 21)
    hrng = linspace(0, 0.1, 21)
    fs, hs = meshgrid(frng, hrng)
    ks = zeros(fs.shape)
    for fi, f in enumerate(frng):
        for hi, h in enumerate(hrng):
            ks[fi, hi] = maxk(human, (f, 0, h), eta)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(fs, hs, ks)
    plt.show()

    # print maxkstable(human, robot)
    # print maxksafe(human, robot, eta)
    print maxk(human, robot, eta)

    print 
    print "***"
    print 

    print bestrobotstable(human)
    print bestrobotsafe(human, eta)
    # print bestrobot(human, eta)
    '''