open Common

type t =
  | Int
  | Bool
  | Float
  | Any
  | Arrow of t list * t
  | Product of t list
  | Vec of t
  | Variant of case list
  | Rec of string * t
  | Var of string

and case = Case of string * t list

let rec pp fmt = function
  | Int -> Format.fprintf fmt "int"
  | Bool -> Format.fprintf fmt "bool"
  | Float -> Format.fprintf fmt "float"
  | Any -> Format.fprintf fmt "?"
  | Arrow (tys, ty) ->
      Format.fprintf fmt "(%a -> %a)" (Pp.pp_list ~sep:" -> " pp) tys pp ty
  | Product l -> Format.fprintf fmt "(%a)" (Pp.pp_list ~sep:" * " pp) l
  | Vec t -> Format.fprintf fmt "vec[%a]" pp t
  | Variant cases ->
      Format.fprintf fmt "(%a)" (Pp.pp_list ~sep:" | " pp_case) cases
  | Rec (type_name, _) -> Format.fprintf fmt "%s" type_name
  | Var type_name -> Format.fprintf fmt "%s" type_name

and pp_case fmt = function
  | Case (name, tys) ->
      if tys = [] then Format.fprintf fmt "%s" name
      else Format.fprintf fmt "%s of %a" name (Pp.pp_list ~sep:" * " pp) tys

let ty_of_op1 : Prim.op1 -> t * t = function
  | Neg -> (Int, Int)
  | Not -> (Bool, Bool)
  | Print print -> (
      match print with
      | PrintInt -> (Int, Int)
      | PrintBool -> (Bool, Bool)
      | PrintFloat -> (Float, Float)
      | PrintAny -> (Any, Any)
      | PrintChar _ -> (Int, Int))
  | ReadInt -> (Int, Int)
  | ReadFloat -> (Int, Float)
  | ReadChar -> (Int, Int)
  | Time -> (Arrow ([ Int ], Int), Int)
  | FNeg -> (Float, Float)
  | FSqrt -> (Float, Float)
  | FSin -> (Float, Float)
  | FCos -> (Float, Float)
  | FExp -> (Float, Float)
  | FLog -> (Float, Float)
  | FRound -> (Float, Float)
  | Float_of_int -> (Int, Float)
  | Int_of_float -> (Float, Int)

let ty_of_op2 : Prim.op2 -> t * t * t = function
  | Plus | Minus | Times | Div | Mod -> (Int, Int, Int)
  | FPlus | FMinus | FTimes | FDiv | FMax | FMin -> (Float, Float, Float)
  | And | Or -> (Bool, Bool, Bool)
  | Eq | Less | LessEq | Greater | GreaterEq -> (Int, Int, Bool)
  | FEq | FLess | FLessEq | FGreater | FGreaterEq -> (Float, Float, Bool)

let dom = function Arrow (a, _) -> Some a | Any -> Some [ Any ] | _ -> None
let cod = function Arrow (_, b) -> Some b | Any -> Some Any | _ -> None

let variant_id =
  let hashtable = Hashtbl.create 10 in
  fun ty ->
    match ty with
    | Rec (_, Variant _) ->
        if Hashtbl.mem hashtable ty then Hashtbl.find hashtable ty
        else
          let id = Hashtbl.length hashtable in
          Hashtbl.add hashtable ty id;
          id
    | _ -> failwith "tried to assign variant id to type that is not a variant"

let case_id casename ty =
  let cases =
    match ty with
    | Rec (_, Variant cases) -> cases
    | _ -> failwith "tried to obtain case id from type that is not a variant"
  in
  let i =
    cases
    |> List.mapi (fun i case -> (i, case))
    |> List.find (fun (_, Case (name, _)) -> name = casename)
    |> fst
  in
  i

let rec subst_ty ty_name ty_subst ty =
  match ty with
  | Int | Bool | Float | Any -> ty
  | Var name -> if name = ty_name then ty_subst else ty
  | Arrow (tys, ty) ->
      Arrow
        (List.map (subst_ty ty_name ty_subst) tys, subst_ty ty_name ty_subst ty)
  | Product ts -> Product (List.map (subst_ty ty_name ty_subst) ts)
  | Vec t -> Vec (subst_ty ty_name ty_subst t)
  | Variant cases ->
      Variant
        (List.map
           (fun (Case (name, tys)) ->
             Case (name, List.map (subst_ty ty_name ty_subst) tys))
           cases)
  | Rec (name, ty) ->
      if name = ty_name then ty else Rec (name, subst_ty ty_name ty_subst ty)
