(* ==================================================================
   Cours 4 (suite) : Encodeur (dense + activation)
   ================================================================== *)

(* ------------------------------------------------------------------
   Un encodeur transforme des donnees en une representation latente
   (code) de dimension plus faible.  On construit les briques :
     - couche lineaire  : y = W*x + b
     - fonctions d'activation : sigmoid, tanh, relu
     - empilement en reseau
   ------------------------------------------------------------------ *)

(* ------------------------------------------------------------------
   1. Fonctions d'activation
   ------------------------------------------------------------------ *)

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.

(* Derivees pour le futur autoencodeur *)
let dsigmoid x =
  let s = sigmoid x in s *. (1. -. s)
let dtanh x =
  let t = tanh x in 1. -. t *. t
let drelu x = if x > 0. then 1. else 0.

(* ------------------------------------------------------------------
   2. Couche lineaire
   ------------------------------------------------------------------ *)

(* Un poids initialise aleatoirement *)
let rand_poids () = Random.float 2. -. 1.

(* Forward d'une couche : y_j = b_j + sum_i W_{j,i} * x_i *)
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

(* ------------------------------------------------------------------
   3. Reseau = pile de couches
   ------------------------------------------------------------------ *)

(* Un reseau est une liste de (poids, biais, activation) *)
type activation = Sigmoid | Tanh | Relu | Linear

type couche = {
    poids : float list list;   (* [sortie][entree] *)
    biais : float list;
    actif : activation;
  }

type reseau = couche list

(* Forward complet du reseau *)
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

(* ------------------------------------------------------------------
   4. Initialisation de 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 }

(* Encodeur : 2 entrees -> 3 cachee (tanh) -> 1 code (linear) *)
let encodeur_2d_1d () =
  [ cree_couche 2 3 Tanh;
    cree_couche 3 1 Linear;
  ]

(* Encodeur plus profond : 2 -> 4 -> 2 -> 1 *)
let encodeur_profond () =
  [ cree_couche 2 4 Relu;
    cree_couche 4 2 Relu;
    cree_couche 2 1 Linear;
  ]

(* ------------------------------------------------------------------
   5. Visualisation du code latent
   ------------------------------------------------------------------ *)

(* Nuage de points en spirale (donnees non lineaires) *)
let genere_spirale n bruit =
  let rec aux i acc =
    if i >= n then acc
    else
      let t = float_of_int i *. 0.3 in
      let r = float_of_int i *. 0.04 in
      let x = r *. cos t +. Random.float bruit -. bruit /. 2.
      and y = r *. sin t +. Random.float bruit -. bruit /. 2. in
      aux (i + 1) ((x, y) :: acc)
  in
  aux 0 []

(* Nuage de points 2D avec 4 clusters *)
let genere_4clusters n_par_cluster =
  let centres = [(-2., -2.); (-2., 2.); (2., -2.); (2., 2.)] in
  List.concat_map (fun (cx, cy) ->
    List.init n_par_cluster (fun _ ->
      cx +. Random.float 2. -. 1., cy +. Random.float 2. -. 1.)
  ) centres

(* Encode un nuage de points et affiche les codes resultants *)
let encode_et_affiche reseau points titre =
  Printf.printf "%s\n" titre;
  Printf.printf "  %-8s  %-8s  ->  code\n" "x" "y";
  let codes = List.map (fun (x, y) ->
    let code = forward reseau [x; y] in
    (x, y, List.hd code)
  ) points in
  (* Trier par code pour voir la structure *)
  let codes = List.sort (fun (_, _, c1) (_, _, c2) -> compare c1 c2) codes in
  List.iter (fun (x, y, code) ->
    Printf.printf "  %6.2f  %6.2f  ->  %+7.4f\n" x y code
  ) codes

(* ------------------------------------------------------------------
   6. Exemples
   ------------------------------------------------------------------ *)

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

  (* --- Exemple 1 : encodeur aleatoire sur 4 clusters --- *)
  Printf.printf "--- Exemple 1 : 4 clusters encodes (poids aleatoires) ---\n";
  let points = genere_4clusters 5 in
  let enc = encodeur_2d_1d () in
  encode_et_affiche enc points "  Avec un encodeur aleatoire :\n";
  Printf.printf "\n  Les codes sont < 0 pour certains clusters, > 0 pour\n";
  Printf.printf "  d'autres : l'encodeur separe deja partiellement les clusters.\n";

  (* --- Exemple 2 : encodeur profond sur spirale --- *)
  Printf.printf "\n--- Exemple 2 : Spirale -> code 1D (encodeur profond) ---\n";
  let points = genere_spirale 20 0.05 in
  let enc = encodeur_profond () in
  encode_et_affiche enc points "\n  Encodeur 2->4->2->1 (Relu), poids aleatoires :\n";
  Printf.printf "\n  Meme avec des poids aleatoires, le code capture\n";
  Printf.printf "  partiellement la structure de la spirale.\n"

(* ------------------------------------------------------------------
   Exercices pour le cours
   ------------------------------------------------------------------ *
   1. Forward par batch : encoder tout un tableau d'un coup (matrice)

   2. Normalisation des entrees (mean = 0, std = 1) avant encodage

   3. Visualisation du code en 2D : modifier l'encodeur pour
      produire un code a 2 dimensions, afficher les points colores
      par cluster dans l'espace latent

   4. Encodeur convolutif : pour des images (tableaux 2D), utiliser
      une convolution comme premiere couche

   5. Initialisation Xavier/Glorot des poids (diviser par sqrt(n_entrees))
   *)