(* ==================================================================
   Cours 4 (suite) : Decodeur (du code latent vers les donnees)
   ================================================================== *)

(* ------------------------------------------------------------------
   Le decodeur fait l'inverse de l'encodeur : il prend un code latent
   et reconstruit les donnees originales.
   ------------------------------------------------------------------ *)

(* ------------------------------------------------------------------
   1. Briques du decodeur (reprises de encoder.ml)
   ------------------------------------------------------------------ *)

let sigmoid x = 1. /. (1. +. exp (-.x))
let tanh x = (exp (2. *. x) -. 1.) /. (exp (2. *. x) +. 1.)
let relu x = if x > 0. then x else 0.

let rand_poids () = Random.float 2. -. 1.

let forward_couche (w : float list list) (b : float list) (x : float list) =
  List.map2 (fun ligne bias ->
    List.fold_left2 (fun s w_ji x_i -> s +. w_ji *. x_i) bias ligne x
  ) w b

type activation = Sigmoid | Tanh | Relu | Linear

type couche = {
    poids : float list list;
    biais : float list;
    actif : activation;
  }

type reseau = couche list

let forward reseau entree =
  List.fold_left (fun x couche ->
    let z = forward_couche couche.poids couche.biais x in
    match couche.actif with
    | Sigmoid -> List.map sigmoid z
    | Tanh    -> List.map tanh z
    | Relu    -> List.map relu z
    | Linear  -> z
  ) entree reseau

let cree_couche entrees sorties actif =
  let poids = List.init sorties (fun _ ->
    List.init entrees (fun _ -> rand_poids ()))
  and biais = List.init sorties (fun _ -> 0.) in
  { poids; biais; actif }

(* ------------------------------------------------------------------
   2. Decodeur : architecture symetrique de l'encodeur
   ------------------------------------------------------------------ *)

(* Decodeur : 1 code -> 3 cachee (tanh) -> 2 sorties (linear) *)
let decodeur_1d_2d () =
  [ cree_couche 1 3 Tanh;
    cree_couche 3 2 Linear;
  ]

(* Decodeur profond : 1 -> 2 -> 4 -> 2 *)
let decodeur_profond () =
  [ cree_couche 1 2 Relu;
    cree_couche 2 4 Relu;
    cree_couche 4 2 Linear;
  ]

(* ------------------------------------------------------------------
   3. Erreur de reconstruction
   ------------------------------------------------------------------ *)

let erreur_reconstruction decodeur code original =
  let reconstruit = forward decodeur code in
  List.fold_left2 (fun s xr xo -> s +. (xr -. xo) ** 2.) 0. reconstruit original

let erreur_moyenne decodeur code_et_original =
  let n = float_of_int (List.length code_et_original) in
  List.fold_left (fun s (code, orig) ->
    s +. erreur_reconstruction decodeur code orig
  ) 0. code_et_original /. n

(* ------------------------------------------------------------------
   4. Decodage et visualisation
   ------------------------------------------------------------------ *)

let decode_et_affiche decodeur codes_originaux titre =
  Printf.printf "%s\n" titre;
  Printf.printf "  %-8s  ->  %-8s  %-8s  (reconstruit)  err\n" "code" "x" "y";
  List.iter (fun (code, (xo, yo)) ->
    let r = forward decodeur [code] in
    let xr = List.nth r 0 and yr = List.nth r 1 in
    let err = sqrt ((xr -. xo) ** 2. +. (yr -. yo) ** 2.) in
    Printf.printf "  %+6.3f  ->  %6.2f  %6.2f  (%6.2f  %6.2f)  %.4f\n"
      code xo yo xr yr err
  ) codes_originaux

(* ------------------------------------------------------------------
   5. Exemples
   ------------------------------------------------------------------ *)

let () =
  Random.self_init ();
  Printf.printf "=== Decodeur ===\n\n";
  Printf.printf "  Le decodeur transforme un code 1D en donnees 2D.\n";
  Printf.printf "  Architecture : 1 -> 3 (tanh) -> 2 (lineaire)\n\n";

  (* --- Exemple 1 : decoder des codes aleatoires --- *)
  Printf.printf "--- Exemple 1 : Codes uniformes -> points 2D ---\n";
  let dec = decodeur_1d_2d () in
  let codes = List.init 8 (fun i ->
    let code = float_of_int (i - 4) *. 0.5 in
    (code, (0., 0.))  (* pas d'original connu *)
  ) in
  Printf.printf "  Decodage de codes regulierement espaces (poids aleatoires) :\n";
  List.iter (fun (code, _) ->
    let r = forward dec [code] in
    Printf.printf "  code %+5.2f  ->  (%.3f, %.3f)\n" code
      (List.nth r 0) (List.nth r 1)
  ) codes;

  (* --- Exemple 2 : Decodeur profond sur la spirale --- *)
  Printf.printf "\n--- Exemple 2 : Spirale encodee puis decodee ---\n";

  (* On encode la spirale avec l'encodeur, puis on decode *)
  let enc = [
    cree_couche 2 3 Tanh;
    cree_couche 3 1 Linear;
  ] in
  let dec = decodeur_1d_2d () in

  (* Points de la spirale *)
  let rec spirale n acc =
    if n <= 0 then acc
    else
      let t = float_of_int n *. 0.3 in
      let r = float_of_int n *. 0.04 in
      spirale (n - 1) ((r *. cos t, r *. sin t) :: acc)
  in
  let points = spirale 15 [] in

  (* Encoder -> Decoder *)
  let codes_originaux = List.map (fun (x, y) ->
    let code = forward enc [x; y] in
    (List.hd code, (x, y))
  ) points in
  decode_et_affiche dec codes_originaux
    "\n  Encode -> Decode (poids tous aleatoires) :\n";
  Printf.printf "\n  Sans entrainement, la reconstruction est mauvaise.\n";
  Printf.printf "  L'autoencodeur (prochaine etape) ajustera les poids\n";
  Printf.printf "  pour minimiser cette erreur de reconstruction.\n";

  (* --- Exemple 3 : Decoder un chemin continu dans le code latent --- *)
  Printf.printf "\n--- Exemple 3 : Interpolation dans le code latent ---\n";
  Printf.printf "  On fait varier le code de -2 a +2 et on observe\n";
  Printf.printf "  la forme du decodage (avec decodeur profond) :\n";
  let dec = decodeur_profond () in
  for i = 0 to 10 do
    let code = -2. +. float_of_int i *. 0.4 in
    let r = forward dec [code] in
    Printf.printf "  code %+5.2f  ->  (%7.4f, %7.4f)\n" code
      (List.nth r 0) (List.nth r 1)
  done;
  Printf.printf "\n  Meme avec des poids aleatoires, la sortie varie\n";
  Printf.printf "  continument avec le code : le decodeur definit\n";
  Printf.printf "  une surface parametree par le code latent.\n"

(* ------------------------------------------------------------------
   Exercices pour le cours
   ------------------------------------------------------------------ *
   1. Inverser la sortie du decodeur avec sigmoid final pour des
      donnees normalisees entre 0 et 1.

   2. Decodeur conditionnel : ajouter une entree supplementaire
      (label) pour controler la classe du point genere.

   3. Visualiser l'espace latent : decoder une grille reguliere
      de codes pour voir la "carte" generee par le decodeur.

   4. Decodeur avec dropout : desactiver aleatoirement des neurones
      pendant l'entrainement pour regulariser.

   5. Mesure de qualite : rapport signal/bruit (PSNR) entre
      l'original et la reconstruction.
   *)