(* ==================================================================
   Cours 4 (suite) : Reseau multi-couche et autoencodeur
   ================================================================== *)

#directory "/home/thiry/Bureau/ML"
#load "graphiques.cmo"
open Graphiques

(* ------------------------------------------------------------------
   On implemente la retropropagation (backpropagation) pour un
   reseau multi-couche.  On l'utilise pour entrainer un autoencodeur :
   encodeur 2D -> 1D -> decodeur 1D -> 2D.
   ------------------------------------------------------------------ *)

(* ------------------------------------------------------------------
   1. Fonctions d'activation et leurs derivees
   ------------------------------------------------------------------ *)

let sigmoid x = 1. /. (1. +. exp (-.x))
let dsigmoid x = let s = sigmoid x in s *. (1. -. s)

let tanh x = (exp (2. *. x) -. 1.) /. (exp (2. *. x) +. 1.)
let dtanh x = let t = tanh x in 1. -. t *. t

let relu x = if x > 0. then x else 0.
let drelu x = if x > 0. then 1. else 0.

(* ------------------------------------------------------------------
   2. Types du reseau
   ------------------------------------------------------------------ *)

type activation = Sigmoid | Tanh | Relu | Linear

type couche = {
    mutable poids : float list list;  (* [sorties][entrees] *)
    mutable biais : float list;
    actif : activation;
  }

type reseau = couche list

(* ------------------------------------------------------------------
   3. Initialisation
   ------------------------------------------------------------------ *)

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

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

(* ------------------------------------------------------------------
   4. Forward avec sauvegarde des activations
   ------------------------------------------------------------------ *)

(* forward_entrainement : calcule les sorties ET sauvegarde
   les entrees/sorties de chaque couche pour la retropropagation *)
let forward_entrainement reseau entree =
  let _ = entree in  (* entrees du reseau *)
  List.fold_left (fun (x, couches_act) couche ->
    let z = List.map2 (fun ligne bias ->
      List.fold_left2 (fun s w_ji x_i -> s +. w_ji *. x_i) bias ligne x
    ) couche.poids couche.biais in
    let a = match couche.actif with
      | Sigmoid -> List.map sigmoid z
      | Tanh    -> List.map tanh z
      | Relu    -> List.map relu z
      | Linear  -> z
    in
    (a, (x, z, a) :: couches_act)
  ) (entree, []) reseau

(* ------------------------------------------------------------------
   5. Retropropagation (backward)
   ------------------------------------------------------------------ *)

let derivee_actif = function
  | Sigmoid -> dsigmoid
  | Tanh    -> dtanh
  | Relu    -> drelu
  | Linear  -> fun _ -> 1.

(* backward : calcule les gradients pour tout le reseau *)
let backward reseau couches_act sortie cible =
  (* Erreur sur la sortie : dE/da = a - cible (MSE) *)
  let da = List.map2 (-.) sortie cible in
  (* Parcours des couches en sens inverse *)
  let (_ , grads_poids, grads_biais) =
    List.fold_right (fun (couche, (x, z, a)) (da_suiv, gp, gb) ->
      let dact = derivee_actif couche.actif in
      (* dE/dz = dE/da * dact(z) *)
      let dz = List.map2 (fun da_i z_i -> da_i *. dact z_i) da_suiv z in
      (* dE/dW = dz * x^T  (produit externe) *)
      let gp_couche = List.map (fun dz_j ->
        List.map (fun x_i -> dz_j *. x_i) x
      ) dz in
      (* dE/db = dz *)
      let gb_couche = dz in
      (* dE/dx = W^T * dz  (pour la couche precedente) *)
      let n_entrees = List.length (List.hd couche.poids) in
      let da_prec = List.init n_entrees (fun i ->
        List.fold_left2 (fun s w_j_i dz_j -> s +. w_j_i *. dz_j)
          0. (List.map (fun ligne -> List.nth ligne i) couche.poids) dz
      ) in
      (da_prec, gp_couche :: gp, gb_couche :: gb)
    ) (List.combine reseau (List.rev couches_act)) (da, [], [])
  in
  (grads_poids, grads_biais)

(* ------------------------------------------------------------------
   6. Mise a jour des poids (SGD)
   ------------------------------------------------------------------ *)

let mise_a_jour reseau grads_poids grads_biais taux =
  List.iter2 (fun couche (gp, gb) ->
    couche.poids <- List.map2 (fun ligne dW ->
      List.map2 (fun w dw -> w -. taux *. dw) ligne dW
    ) couche.poids gp;
    couche.biais <- List.map2 (fun b db -> b -. taux *. db) couche.biais gb
  ) reseau (List.combine grads_poids grads_biais)

(* ------------------------------------------------------------------
   7. Entrainement complet (SGD)
   ------------------------------------------------------------------ *)

let entraine reseau donnees taux n_iter =
  let n = float_of_int (List.length donnees) in
  for iter = 1 to n_iter do
    let cout_total = ref 0. in
    (* On traite chaque exemple un par un (SGD en ligne) *)
    List.iter (fun (entree, cible) ->
      let (sortie, couches_act) = forward_entrainement reseau entree in
      let (gp, gb) = backward reseau couches_act sortie cible in
      mise_a_jour reseau gp gb (taux /. n);
      (* Erreur quadratique *)
      cout_total := !cout_total +.
        List.fold_left2 (fun s a b -> s +. (a -. b) ** 2.) 0. sortie cible
    ) donnees;
    if iter mod 100 = 0 || iter = n_iter then
      Printf.printf "  iter %4d : cout = %.6f\n%!" iter (!cout_total /. n)
  done

(* ------------------------------------------------------------------
   8. Construction de l'autoencodeur
   ------------------------------------------------------------------ *)

(* Autoencodeur : 2 -> 3 (tanh) -> 1 (linear) -> 3 (tanh) -> 2 (linear) *)
let autoencodeur () =
  [ cree_couche 2 3 Tanh;    (* encodeur *)
    cree_couche 3 1 Linear;
    cree_couche 1 3 Tanh;    (* decodeur *)
    cree_couche 3 2 Linear;
  ]

(* Autoencodeur profond : 2 -> 4 -> 2 -> 1 -> 2 -> 4 -> 2 *)
let autoencodeur_profond () =
  [ cree_couche 2 4 Tanh;
    cree_couche 4 2 Tanh;
    cree_couche 2 1 Linear;
    cree_couche 1 2 Tanh;
    cree_couche 2 4 Tanh;
    cree_couche 4 2 Linear;
  ]

(* ------------------------------------------------------------------
   9. Extraction encodeur / decodeur depuis l'autoencodeur
   ------------------------------------------------------------------ *)

(* L'autoencodeur est [enc1; enc2; dec1; dec2]
   L'encodeur seul = [enc1; enc2]
   Le decodeur seul = [dec1; dec2] *)

let extrait_encodeur ae = [List.nth ae 0; List.nth ae 1]
let extrait_decodeur ae = [List.nth ae 2; List.nth ae 3]

let forward_fwd reseau entree =
  fst (forward_entrainement reseau entree)

(* ------------------------------------------------------------------
   10. Generation de donnees : 4 clusters
   ------------------------------------------------------------------ *)

let genere_donnees n_par_cluster =
  let centres = [(-2., -2., 0.); (-2., 2., 1.); (2., -2., 2.); (2., 2., 3.)] in
  List.concat_map (fun (cx, cy, _) ->
    List.init n_par_cluster (fun _ ->
      let x = cx +. Random.float 2. -. 1.
      and y = cy +. Random.float 2. -. 1. in
      ([x; y], [x; y])  (* entree = cible : autoencodeur *)
    )
  ) centres

(* ------------------------------------------------------------------
   11. Evaluation de la reconstruction
   ------------------------------------------------------------------ *)

let erreur_moyenne reseau donnees =
  let n = float_of_int (List.length donnees) in
  List.fold_left (fun s (entree, cible) ->
    let sortie = forward_fwd reseau entree in
    s +. List.fold_left2 (fun s2 a b -> s2 +. (a -. b) ** 2.) 0. sortie cible
  ) 0. donnees /. n

(* ------------------------------------------------------------------
   12. Exemples
   ------------------------------------------------------------------ *)

let () =
  Random.self_init ();
  Printf.printf "=== Autoencodeur (reseau multi-couche) ===\n\n";
  Printf.printf "  Architecture : 2 -> 3 (tanh) -> 1 (linear)\n";
  Printf.printf "                -> 3 (tanh) -> 2 (linear)\n\n";

  let donnees = genere_donnees 40 in
  let n_train = List.length donnees * 8 / 10 in
  let train = List.filteri (fun i _ -> i < n_train) donnees
  and test  = List.filteri (fun i _ -> i >= n_train) donnees in

  (* --- Exemple 1 : avant entrainement --- *)
  Printf.printf "--- Exemple 1 : Reconstruction AVANT entrainement ---\n";
  let ae = autoencodeur () in
  Printf.printf "  Erreur moyenne (train) : %.6f\n" (erreur_moyenne ae train);
  Printf.printf "  Erreur moyenne (test)  : %.6f\n" (erreur_moyenne ae test);
  Printf.printf "\n  Quelques reconstructions (entree -> code -> sortie) :\n";
  let quelques_uns = List.filteri (fun i _ -> i mod 20 = 0) train in
  List.iter (fun (entree, cible) ->
    let sortie = forward_fwd ae entree in
    let code = forward_fwd (extrait_encodeur ae) entree in
    Printf.printf "  (%.2f, %.2f) -> [%+.4f] -> (%.2f, %.2f)\n"
      (List.nth entree 0) (List.nth entree 1)
      (List.hd code)
      (List.nth sortie 0) (List.nth sortie 1)
  ) quelques_uns;
  (* --- Exemple 2 : apres entrainement --- *)
  Printf.printf "\n--- Exemple 2 : Entrainement en cours... ---\n";
  let ae = autoencodeur () in
  entraine ae train 0.5 1000;

  Printf.printf "\n--- Exemple 3 : Reconstruction APRES entrainement ---\n";
  Printf.printf "  Erreur moyenne (train) : %.6f\n" (erreur_moyenne ae train);
  Printf.printf "  Erreur moyenne (test)  : %.6f\n" (erreur_moyenne ae test);
  Printf.printf "\n  Reconstructions apres entrainement :\n";
  let quelques_uns = List.filteri (fun i _ -> i mod 20 = 0) test in
  List.iter (fun (entree, cible) ->
    let sortie = forward_fwd ae entree in
    let code = forward_fwd (extrait_encodeur ae) entree in
    Printf.printf "  (%.2f, %.2f) -> [%+.4f] -> (%.2f, %.2f)\n"
      (List.nth entree 0) (List.nth entree 1)
      (List.hd code)
      (List.nth sortie 0) (List.nth sortie 1)
  ) quelques_uns;

  Printf.printf "\n  Le code latent de quelques points (separes par cluster) :\n";
  let tous_codes = List.map (fun (e, _) ->
    forward_fwd (extrait_encodeur ae) e
  ) test in
  List.iteri (fun i c ->
    Printf.printf "  point %2d : code = %+.4f\n" i (List.hd c)
  ) tous_codes

(* ------------------------------------------------------------------
   Graphiques SVG
   ------------------------------------------------------------------ *)

let () =
  let ae = autoencodeur () in
  let donnees = genere_donnees 40 in
  let n = float_of_int (List.length donnees) in
  let historique = ref [] in
  for iter = 1 to 1000 do
    let cout_total = ref 0. in
    List.iter (fun (entree, cible) ->
      let (sortie, couches_act) = forward_entrainement ae entree in
      let (gp, gb) = backward ae couches_act sortie cible in
      mise_a_jour ae gp gb (0.5 /. n);
      cout_total := !cout_total +.
        List.fold_left2 (fun s a b -> s +. (a -. b) ** 2.) 0. sortie cible
    ) donnees;
    if iter mod 20 = 0 then
      historique := (float_of_int iter, !cout_total /. n) :: !historique
  done;

  courbe_entrainement ~fichier:"autoencodeur_loss.svg"
    ~titre:"Entrainement de l'autoencodeur"
    ~xlab:"Iteration" ~ylab:"Erreur quadratique"
    [{ points = List.rev !historique; legende = "MSE" }]

let () =
  let ae = autoencodeur () in
  let donnees = genere_donnees 40 in
  entraine ae donnees 0.5 1000;
  (* Codes latents des donnees de test *)
  let codes = List.map (fun (e, _) ->
    let c = forward_fwd (extrait_encodeur ae) e in
    (List.nth e 0, List.hd c)
  ) donnees in
  let codes_par_idx = [List.filteri (fun i _ -> i mod 2 = 0) codes;
                       List.filteri (fun i _ -> i mod 2 = 1) codes] in
  nuage ~fichier:"autoencodeur_codes.svg"
    ~titre:"Code latent 1D (en fonction de x)"
    ~xlab:"x" ~ylab:"valeur du code"
    codes_par_idx

(* ------------------------------------------------------------------
   Exercices pour le cours
   ------------------------------------------------------------------ *
   1. Visualiser le code latent 1D pour les 4 clusters : les points
      d'un meme cluster ont-ils des codes proches ?

   2. Autoencodeur denoissant : ajouter du bruit a l'entree mais
      garder la sortie propre.

   3. Autoencodeur contractif : ajouter une penalite sur la norme
      du Jacobien de l'encodeur.

   4. Variational Autoencoder (VAE) : remplacer le code latent
      deterministe par une distribution (moyenne + variance).

   5. Autoencodeur pour des images : utiliser le meme code avec
      des entrees 8x8 (64 pixels) -> 16 -> 2 -> 16 -> 64.

   6. Early stopping : arreter l'entrainement quand l'erreur de
      test cesse de diminuer.
   *)