import pygame
from pygame.locals import KEYDOWN, K_SPACE, K_w, K_s, K_a, K_d,K_ESCAPE,K_SPACE, QUIT, K_LEFT, K_RIGHT
import numpy as np


import Box2D  # The main library
# Box2D.b2 maps Box2D.b2Vec2 to vec2 (and so on)
from Box2D.b2 import (world, polygonShape, staticBody, dynamicBody,circleShape)
from Box2D import (b2FixtureDef, b2PolygonShape, b2Vec2)

PPM = 15
TARGET_FPS = 60
TIME_STEP = 1.0/TARGET_FPS
SCREEN_WIDTH, SCREEN_HEIGHT = 640 , 480

screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT), 0, 32)
pygame.display.set_caption('Top-Down Car')
clock = pygame.time.Clock()

# Create the world
GRAVITY = [0, 0]
world = world(gravity=(GRAVITY), doSleep=True)

Top_Down_BodyWork_Shape = [(11.5, 0.0),
            (13.0, 2.5),
            (12.25, 5.5),
            (11, 7.5),
            (10.25, 10.0),
            (9.75, 10.0),
            (9, 7.5),
            (7.75, 5.5),
            (7.0, 2.5),
            (8.5, 0.0),
            ] # 1 Axis of symmetry

Tire_Anchoring = [(7.0, 0.75),
            (13.0, 0.75),
            (7.0, 8.50),
            (13.0, 8.50),
            ]

### TOP DOWN CAR DEFINITION ###

Wheels = []
Suspensions = []

Top_Down_BodyWork = world.CreateDynamicBody(position=(0, 0),
                                        angle=0,
                                        fixtures = b2FixtureDef(
                                            shape = b2PolygonShape(vertices=Top_Down_BodyWork_Shape),
                                            density = 2,
                                            friction = 0.3))

for i in range(4):
# Create Wheels
Wheels.append(world.CreateDynamicBody(position=(Tire_Anchoring[i]),
                                    angle=0,
                                    fixtures=b2FixtureDef(
                                            shape=b2PolygonShape(box=(0.5, 1.5)),
                                            density=1.0,
                                            friction = 1)))

# Anchor Wheels to Bodywork

if i in range(2):
    # Rear Wheels do not move so they can be welded to bodywork
    Suspensions.append(world.CreateWeldJoint(bodyA=Top_Down_BodyWork,
                                            bodyB= Wheels[i],
                                            localAnchorA=Tire_Anchoring[i],
                                            localAnchorB=(0,0),
                                            ))
else:
    # Front wheels are revolute joints
    Suspensions.append(world.CreateRevoluteJoint(bodyA=Top_Down_BodyWork,
                                            bodyB= Wheels[i],
                                            localAnchorA=Tire_Anchoring[i],
                                            localAnchorB=(0,0),
                                            enableMotor=True,
                                            maxMotorTorque=100000,
                                            enableLimit=True,
                                            lowerAngle= -np.pi/4,
                                            upperAngle= np.pi/4,
                                            motorSpeed = 0,
                                            ))

Left_Wheel, Right_Wheel = Wheels[2], Wheels[3]
Left_Susp, Right_Susp = Suspensions[2], Suspensions[3]

### DRAWING ###

colors = {dynamicBody: (133, 187, 101, 0),  staticBody: (15, 0, 89, 0)}


def my_draw_polygon(polygon, body, fixture):
vertices = [(body.transform * v) * PPM for v in polygon.vertices]
vertices = [(v[0], SCREEN_HEIGHT - v[1]) for v in vertices]
pygame.draw.polygon(screen, colors[body.type], vertices)
polygonShape.draw = my_draw_polygon


def my_draw_circle(circle, body, fixture):
position = body.transform * circle.pos * PPM
position = (position[0], SCREEN_HEIGHT - position[1])
pygame.draw.circle(screen, colors[body.type], [int(
    x) for x in position], int(circle.radius * PPM))
circleShape.draw= my_draw_circle

### MAIN GAME LOOP ###

SteerAngle = 0
Speed = 0
dV = 0.01 #Speed Increment

running = True
number_of_states = 0
max_lateral_impulse = 30
Right_Normal_Squared = []
tick_counter = 0

while running:
# Check the event queue
for event in pygame.event.get():
    if event.type == QUIT or (event.type == KEYDOWN and event.key == K_ESCAPE):
        # The user closed the window or pressed escape
        running = False

screen.fill((255, 255, 255,0))
for body in world.bodies:
    for fixture in body.fixtures:
        fixture.shape.draw(body, fixture)

    # Cancelling lateral velocity
    for wheel in Wheels:

        if wheel in range(2):
            continue
        else:


            Current_Normal = wheel.GetWorldVector((0,1))
            Current_N_Squared = Current_Normal.dot(wheel.linearVelocity) * Current_Normal

            Right_Normal = wheel.GetWorldVector((1,0))
            Right_Normal_Squared = Right_Normal.dot(wheel.linearVelocity) * Right_Normal

            Impulse = -Right_Normal_Squared * wheel.mass
            if Impulse.length >max_lateral_impulse:
                Impulse *= max_lateral_impulse/Impulse.length

            wheel.ApplyLinearImpulse(Impulse, wheel.worldCenter, True)
            # Allowing skidding
            aimp = 0.25 * wheel.inertia * -wheel.angularVelocity
            wheel.ApplyAngularImpulse(aimp, True)

            current_forward_speed = Current_N_Squared.Normalize()

            drag_force_magnitude = -2 * current_forward_speed
            wheel.ApplyForce(drag_force_magnitude * Current_N_Squared,
                                    wheel.worldCenter, True)

    # Top Down Car Control (WASD keys)

    if event.type == KEYDOWN and event.key == K_w:

        Left_Wheel.ApplyLinearImpulse(b2Vec2(10,0), Left_Wheel.position,False)
        Right_Wheel.ApplyLinearImpulse(b2Vec2(10,0), Left_Wheel.position,False)

    elif event.type == KEYDOWN and event.key == K_s:
        Left_Wheel.ApplyLinearImpulse(b2Vec2(-10,0), Left_Wheel.position,True)
        Right_Wheel.ApplyLinearImpulse(b2Vec2(-10,0), Left_Wheel.position,True)


if event.type == KEYDOWN and event.key==K_a:
        Left_Susp.motorSpeed = 1.5
        Right_Susp.motorSpeed = 1.5

elif event.type == KEYDOWN and event.key == K_d:
    Left_Susp.motorSpeed = -1.5
    Right_Susp.motorSpeed = -1.5

world.Step(TIME_STEP, 10, 10)
tick_counter += 1

if tick_counter > TARGET_FPS:
    Left_Susp.motorSpeed = 0.
    Right_Susp.motorSpeed = 0.
    tick_counter = 0

pygame.display.flip()
clock.tick(TARGET_FPS)

pygame.quit()
print('Done!')