Robot4.elm

module Main exposing (..)

import Exts.Float exposing (roundTo)
import Html exposing (..)
import Html.Attributes exposing (..)
import Svg exposing (Svg)
import Svg.Attributes as SA
import Svg.Path exposing (pathToString, subpath, startAt, lineToMany, emptySubpath, closed, open)
import Mouse exposing (Position)


main : Program Never Model Msg
main =
    Html.program
        { init = init
        , view = view
        , update = update
        , subscriptions = subscriptions
        }



-- MODEL


type alias Robot =
    { origin : Pos
    , upperArmLength : Length
    , forearmLength : Length
    }


type alias Model =
    { robot : Robot
    , armConfigs : ( ConfigPos, ConfigPos )
    , impossible : Bool
    , configSpacePos : Maybe ConfigPos
    , path : ( List ConfigPos, List ConfigPos )
    , pathColor : Color
    , drawingPath : Bool
    , paths : List ( Color, List ConfigPos )
    , pathColors : List Color
    }


type alias HandPosition =
    Pos


type alias ConfigPos =
    ( Angle, Angle )


init : ( Model, Cmd Msg )
init =
    let
        ( w, h ) =
            ( 400, 400 )

        robot =
            { origin = ( w / 2, h / 2 )
            , upperArmLength = 100
            , forearmLength = 75
            }

        armConfigs =
            case inverseKinematics robot ( 100, 100 ) of
                Just solution ->
                    solution

                Nothing ->
                    Debug.crash "No initial solution for robot!"

        model =
            { robot = robot
            , armConfigs = armConfigs
            , impossible = False
            , configSpacePos = Nothing
            , path = ( [], [] )
            , pathColor = "#d523fd"
            , paths = []
            , drawingPath = False
            , pathColors =
                [ "#fd23a2"
                , "#2395fd"
                , "#23fdb0"
                , "#ecfd23"
                , "#fd8b23"
                , "#d523fd"
                ]
            }
    in
        ( model, Cmd.none )



-- UPDATE


type Msg
    = MouseMove Mouse.Position
    | MouseDown Mouse.Position
    | MouseUp Mouse.Position


update : Msg -> Model -> ( Model, Cmd Msg )
update msg model =
    updateHelper msg model ! []


updateHelper : Msg -> Model -> Model
updateHelper msg model =
    case msg of
        MouseMove pos ->
            updatePos pos model
                |> addToPath

        MouseDown pos ->
            model
                |> updatePos pos
                |> startPath
                |> addToPath

        MouseUp pos ->
            model
                |> updatePos pos
                |> addToPath
                |> endPath


updatePos : Mouse.Position -> Model -> Model
updatePos pos model =
    if 0 <= pos.x && pos.x < 400 && 0 <= pos.y && pos.y < 400 then
        updateEuclidean (mouseToHandPos model.robot.origin pos.x pos.y) model
    else if 450 <= pos.x && pos.x <= 750 && 50 <= pos.y && pos.y <= 350 then
        let
            config =
                mouseToConfigSpace (pos.x - 450) (pos.y - 50)
        in
            let
                ( x, y ) =
                    forwardKinematics model.robot config |> .hand
            in
                updateEuclidean (mouseToHandPos model.robot.origin (round x) (round y)) model
    else
        model


startPath : Model -> Model
startPath model =
    { model | drawingPath = True, path = ( [], [] ) }


addToPath : Model -> Model
addToPath model =
    if model.drawingPath && not model.impossible then
        let
            ( p1, p2 ) =
                model.path

            ( c1, c2 ) =
                model.armConfigs
        in
            { model
                | path = ( c1 :: p1, c2 :: p2 )
            }
    else
        model


endPath : Model -> Model
endPath model =
    let
        ( p1, p2 ) =
            model.path

        c =
            model.pathColor

        paths =
            model.paths ++ [ ( c, p1 ), ( c, p2 ) ]
    in
        case model.pathColors of
            [] ->
                { model | drawingPath = False }

            color :: rest ->
                { model
                    | drawingPath = False
                    , pathColor = color
                    , pathColors = rest ++ [ color ]
                    , paths = paths
                }


updateEuclidean : HandPosition -> Model -> Model
updateEuclidean handPos model =
    case inverseKinematics model.robot handPos of
        Nothing ->
            { model | impossible = True }

        Just solutions ->
            { model | impossible = False, armConfigs = solutions }


mouseToHandPos : Pos -> Int -> Int -> HandPosition
mouseToHandPos ( ox, oy ) x y =
    ( toFloat x - ox, toFloat y - oy )


mouseToConfigSpace : Int -> Int -> ConfigPos
mouseToConfigSpace x y =
    let
        px =
            2 * pi * toFloat x / 300 - pi

        py =
            2 * pi * toFloat y / 300 - pi
    in
        ( px, py )



-- SUBSCRIPTIONS


subscriptions : Model -> Sub Msg
subscriptions model =
    Sub.batch
        [ Mouse.moves MouseMove
        , Mouse.downs MouseDown
        , Mouse.ups MouseUp
        ]



-- ROBOT KINEMATICS


type alias Pos =
    ( Float, Float )


type alias Angle =
    Float


type alias Length =
    Float


type alias Positions =
    { shoulder : Pos
    , elbow : Pos
    , hand : Pos
    }


forwardKinematics : Robot -> ConfigPos -> Positions
forwardKinematics robot armConfig =
    let
        ( x0, y0 ) =
            robot.origin

        ( a1, a2 ) =
            armConfig

        ( x1, y1 ) =
            rotate ( x0, y0 ) a1 robot.upperArmLength

        ( x2, y2 ) =
            rotate ( x1, y1 ) (a1 + a2) robot.forearmLength
    in
        { shoulder = ( x0, y0 )
        , elbow = ( x1, y1 )
        , hand = ( x2, y2 )
        }


rotate : Pos -> Angle -> Length -> Pos
rotate ( x0, y0 ) angle len =
    let
        dx =
            len * cos angle

        dy =
            len * sin angle
    in
        ( x0 + dx, y0 + dy )


inverseKinematics : Robot -> Pos -> Maybe ( ConfigPos, ConfigPos )
inverseKinematics robot ( x, y ) =
    let
        l1 =
            robot.upperArmLength

        l2 =
            robot.forearmLength

        r1 =
            x * x + y * y

        r2 =
            l1 * l1 + l2 * l2

        b1 =
            acos ((r1 - r2) / (2 * l1 * l2))

        b2 =
            -b1

        a1 =
            atan2 y x - atan2 (l2 * sin b1) (l1 + l2 * cos b1)

        a2 =
            atan2 y x - atan2 (l2 * sin b2) (l1 + l2 * cos b2)
    in
        if isNaN b1 then
            Nothing
        else
            Just ( ( toRange a1, toRange b1 ), ( toRange a2, toRange b2 ) )


toRange : Float -> Float
toRange angle =
    if angle > pi then
        (angle - 2 * pi)
    else if angle <= -pi then
        (angle + 2 * pi)
    else
        angle



-- VIEW


(=>) : a -> b -> ( a, b )
(=>) =
    (,)


view : Model -> Html Msg
view model =
    div
        [ style
            [ "display" => "flex"
            , "align-items" => "center"
            , "justify-content" => "left"
            ]
        ]
        [ viewRobot model
        , viewConfigSpace model
        ]


listOfArmConfigs : Model -> List ConfigPos
listOfArmConfigs model =
    if model.impossible then
        []
    else
        let
            ( c1, c2 ) =
                model.armConfigs
        in
            [ c1, c2 ]


viewRobot : Model -> Html Msg
viewRobot model =
    let
        fkPath configs =
            List.map (\config -> (forwardKinematics model.robot config).hand) configs

        pathDrawings =
            drawPaths (\p -> forwardKinematics model.robot p |> .hand) (allPaths model)

        robotDrawing =
            drawRobot model.robot (listOfArmConfigs model)

        info =
            showFK model.robot (listOfArmConfigs model)

        annulus =
            drawRangeAnnulus model.robot
    in
        div
            [ style
                [ "background-color" => "#333"
                , "color" => "#fff"
                , "cursor" => "move"
                , "width" => "400px"
                , "height" => "400px"
                , "position" => "relative"
                , "display" => "inline-block"
                ]
            ]
            [ Svg.svg
                [ style
                    [ "position" => "absolute"
                    , "left" => "0"
                    , "top" => "0"
                    , "width" => "400px"
                    , "height" => "400px"
                    ]
                ]
                [ annulus
                , pathDrawings
                , robotDrawing
                ]
            , caption
                [ style
                    [ "position" => "absolute"
                    , "text-align" => "center"
                    , "width" => "100%"
                    , "bottom" => "10px"
                    ]
                ]
                [ text "Euclidean Space" ]
            , pre
                [ style
                    [ "position" => "absolute"
                    , "left" => "0"
                    , "top" => "0"
                    ]
                ]
                [ text info ]
            ]


drawRangeAnnulus : Robot -> Svg Msg
drawRangeAnnulus robot =
    let
        l1 =
            robot.upperArmLength

        l2 =
            robot.forearmLength

        inner =
            abs (l2 - l1)

        outer =
            l2 + l1

        ( ox, oy ) =
            robot.origin

        innerCircle =
            Svg.circle
                [ SA.r (px inner)
                , SA.cx (px ox)
                , SA.cy (px oy)
                , SA.fill "#333"
                ]
                []

        outerCircle =
            Svg.circle
                [ SA.r (px outer)
                , SA.cx (px ox)
                , SA.cy (px oy)
                , SA.fill "#000"
                ]
                []
    in
        Svg.g [] [ outerCircle, innerCircle ]


drawRobot : Robot -> List ConfigPos -> Svg Msg
drawRobot robot configs =
    Svg.g [] (List.map (drawPossibleRobot robot) configs)


drawPossibleRobot : Robot -> ConfigPos -> Svg Msg
drawPossibleRobot robot armConfig =
    let
        { shoulder, elbow, hand } =
            forwardKinematics robot armConfig
    in
        Svg.g []
            [ drawRobotArm shoulder elbow "#3c8d2f"
            , drawRobotArm elbow hand "#2f6a8d"
            , drawRobotHand hand 5 "#8d2f2f"
            ]


showFK : Robot -> List ConfigPos -> String
showFK robot configs =
    case List.head configs of
        Nothing ->
            ""

        Just armConfig ->
            let
                { elbow, hand } =
                    forwardKinematics robot armConfig
            in
                String.join "\n"
                    [ "Elbow position: " ++ showPos elbow
                    , "Hand position: " ++ showPos hand
                    , "Angles: " ++ showConfigRadians armConfig
                    , "Angles: " ++ showConfigDegrees armConfig
                    ]


showConfigRadians : Pos -> String
showConfigRadians ( angle1, angle2 ) =
    "(" ++ showFloat angle1 ++ "rad, " ++ showFloat angle2 ++ "rad)"


showConfigDegrees : Pos -> String
showConfigDegrees ( angle1, angle2 ) =
    "(" ++ showFloat (angle1 * 180 / pi) ++ "°, " ++ showFloat (angle2 * 180 / pi) ++ "°)"


showPos : Pos -> String
showPos ( x, y ) =
    "(" ++ showFloat x ++ ", " ++ showFloat y ++ ")"


showFloat : Float -> String
showFloat =
    toString << roundTo 2


type alias Color =
    String


drawRobotArm : Pos -> Pos -> Color -> Svg msg
drawRobotArm ( x1, y1 ) ( x2, y2 ) color =
    Svg.line
        [ SA.x1 (px x1)
        , SA.y1 (px y1)
        , SA.x2 (px x2)
        , SA.y2 (px y2)
        , SA.stroke color
        , SA.strokeWidth (px 8)
        , SA.strokeLinecap "round"
        ]
        []


drawRobotHand : Pos -> Float -> Color -> Svg msg
drawRobotHand ( x, y ) r color =
    Svg.circle
        [ SA.cx (px x)
        , SA.cy (px y)
        , SA.r (px r)
        , SA.stroke color
        , SA.fill color
        ]
        []


drawPaths : (ConfigPos -> Pos) -> List ( Color, List Pos ) -> Svg msg
drawPaths f paths =
    Svg.g []
        (List.map (\( col, path ) -> drawPath (List.map f path) col) paths)


drawPath : List Pos -> Color -> Svg msg
drawPath path color =
    Svg.path
        [ SA.d (pathToString [ polygon path ])
        , SA.stroke color
        , SA.fill "none"
        , SA.strokeWidth (px 2)
        ]
        []


polygon : List Pos -> Svg.Path.Subpath
polygon ps =
    case ps of
        [] ->
            emptySubpath

        x :: xs ->
            subpath (startAt x) open [ lineToMany xs ]


allPaths : Model -> List ( Color, List ConfigPos )
allPaths model =
    if model.drawingPath then
        let
            ( p1, p2 ) =
                model.path

            c =
                model.pathColor
        in
            ( c, p1 ) :: ( c, p2 ) :: model.paths
    else
        model.paths


viewConfigSpace : Model -> Html Msg
viewConfigSpace model =
    div
        [ style
            [ "display" => "inline-block"
            , "background-color" => "#333"
            , "color" => "white"
            , "width" => "400px"
            , "height" => "400px"
            , "position" => "relative"
            ]
        ]
        [ pre
            [ style
                [ "margin-left" => "50px"
                ]
            ]
            [ text <| showArmConfigs model ]
        , Svg.svg
            [ style
                [ "border" => "2px solid #3c8d2f"
                , "border-left-color" => "#2f6a8d"
                , "border-right-color" => "#2f6a8d"
                , "position" => "absolute"
                , "top" => "50px"
                , "left" => "50px"
                , "width" => "300px"
                , "height" => "300px"
                ]
            ]
            [ drawGrid ( 0, 0 ) ( 300, 300 ) ( 50, 50 )
            , drawPaths (\p -> armConfigToPixels 300 300 p) (allPaths model)
            , drawTargets 300 300 model
            ]
        , caption
            [ style
                [ "position" => "absolute"
                , "text-align" => "center"
                , "width" => "100%"
                , "bottom" => "10px"
                ]
            ]
            [ text "Configuration Space" ]
        , drawAxisCaption ( 45, 355 ) "-π"
        , drawAxisCaption ( 197, 355 ) "0"
        , drawAxisCaption ( 350, 355 ) "π"
        , drawAxisCaption ( 30, 338 ) "-π"
        , drawAxisCaption ( 30, 194 ) "0"
        , drawAxisCaption ( 30, 45 ) "π"
        ]


showArmConfigs : Model -> String
showArmConfigs model =
    listOfArmConfigs model
        |> List.map showConfigRadians
        |> String.join "\n"


drawGrid : ( Float, Float ) -> ( Float, Float ) -> ( Float, Float ) -> Svg Msg
drawGrid ( x0, y0 ) ( x1, y1 ) ( dx, dy ) =
    let
        horizLines =
            List.map (\y -> gridLine ( x0, y ) ( x1, y )) (range (y0 + dx) (y1 - dx) dy)

        vertLines =
            List.map (\x -> gridLine ( x, y0 ) ( x, y1 )) (range (x0 + dx) (x1 - dx) dx)
    in
        Svg.g [] (horizLines ++ vertLines)


gridLine : ( Float, Float ) -> ( Float, Float ) -> Svg Msg
gridLine ( x1, y1 ) ( x2, y2 ) =
    Svg.line
        [ SA.x1 (px x1)
        , SA.y1 (px y1)
        , SA.x2 (px x2)
        , SA.y2 (px y2)
        , SA.stroke "#999"
        , SA.strokeWidth "1px"
        ]
        []


drawAxisCaption : ( Float, Float ) -> String -> Svg Msg
drawAxisCaption ( x, y ) msg =
    caption
        [ style
            [ "position" => "absolute"
            , "text-align" => "center"
            , "left" => px x
            , "top" => px y
            , "color" => "#ddd"
            ]
        ]
        [ text msg ]


range : Float -> Float -> Float -> List Float
range a b d =
    if a > b then
        []
    else
        a :: range (a + d) b d


drawTargets : Float -> Float -> Model -> Svg Msg
drawTargets w h model =
    listOfArmConfigs model
        |> List.map (drawTarget w h)
        |> Svg.g []


drawTarget : Float -> Float -> ConfigPos -> Svg Msg
drawTarget w h ( a1, a2 ) =
    let
        ( cx, cy ) =
            armConfigToPixels w h ( a1, a2 )
    in
        Svg.circle
            [ SA.cx (px cx)
            , SA.cy (px cy)
            , SA.r "3px"
            , SA.stroke "#fff"
            ]
            []


armConfigToPixels : Float -> Float -> ConfigPos -> Pos
armConfigToPixels w h ( a1, a2 ) =
    let
        x =
            (a1 + pi) / (2 * pi)

        y =
            (a2 + pi) / (2 * pi)
    in
        ( w * x, h * y )


px : Float -> String
px number =
    toString number ++ "px"