import do_mpc
from casadi import SX, vertcat
import numpy as np

def create_model():
    from casadi import SX, vertcat

    # Estados do robô: [x, y, theta]
    x = SX.sym('x')        # Posição X
    y = SX.sym('y')        # Posição Y
    theta = SX.sym('theta')  # Orientação (radianos)
    states = vertcat(x, y, theta)

    # Controles do robô: [v, delta]
    v = SX.sym('v')        # Velocidade (m/s)
    delta = SX.sym('delta')  # Ângulo de direção (radianos)
    controls = vertcat(v, delta)

    # Parâmetros do robô
    L = 2.5  # Distância entre eixos (m)

    # Dinâmicas do sistema
    x_dot = v * SX.cos(theta)
    y_dot = v * SX.sin(theta)
    theta_dot = v / L * SX.tan(delta)

    dynamics = vertcat(x_dot, y_dot, theta_dot)

    # Criando o modelo contínuo
    model = do_mpc.model.Model('continuous')

    # Definindo estados, controles e dinâmicas
    model.set_variable(var_type='_x', var_name='x', shape=(3, 1))  # Estados
    model.set_variable(var_type='_u', var_name='u', shape=(2, 1))  # Controles
    model.set_rhs('x', dynamics)  # Define as dinâmicas do sistema

    # Configuração do modelo
    model.setup()
    return model



def create_mpc(model, target_x, target_y, target_theta, inside_street):
    mpc = do_mpc.controller.MPC(model)

    setup_mpc = {
        'n_horizon': 20,
        't_step': 0.2,
        'state_discretization': 'collocation',
        'collocation_type': 'radau',
        'collocation_deg': 3,
        'nlpsol_opts': {'ipopt.print_level': 0, 'ipopt.tol': 1e-4}
    }
    mpc.set_param(**setup_mpc)

    weight_position = 3.0 if inside_street else 10.0
    weight_alignment = 5.0 if inside_street else 7.0

    mpc.set_objective(
        lterm=weight_position * ((model.x['x'][0] - target_x)**2 + (model.x['x'][1] - target_y)**2) +
              weight_alignment * (model.x['x'][2] - target_theta)**2,
        mterm=0
    )

    mpc.bounds['lower', '_u', 'u'] = np.array([0.0, -0.5])
    mpc.bounds['upper', '_u', 'u'] = np.array([3.0, 0.5])

    mpc.setup()
    return mpc


def create_simulator(model):
    simulator = do_mpc.simulator.Simulator(model)
    simulator.set_param(t_step=0.2)
    simulator.setup()
    return simulator


if __name__ == '__main__':
    model = create_model()
    x0 = np.array([0, 0, 0])
    target_x, target_y = 10, 10
    target_theta = np.radians(45)
    inside_street = True

    mpc = create_mpc(model, target_x, target_y, target_theta, inside_street)
    simulator = create_simulator(model)

    simulator.x0 = x0
    mpc.x0 = x0

    for i in range(50):
        u0 = mpc.make_step(x0)
        y_next = simulator.make_step(u0)
        x0 = y_next
        inside_street = np.linalg.norm([x0[0] - target_x, x0[1] - target_y]) < 5.0
        print(f"Passo {i+1}: Estado = {y_next}, Controle = {u0}")