#---- Get environment-specific GVs (env_gvs), their (co)variances (gv_var), covariates and slopes ----

env_gv <- DH@gv[, -1] %*% t(covs_tpe)              # Calculate environment specific genetic values 
env_gv_cent <- scale(env_gv, scale = FALSE)        # Centered the genetic values on zero
env_gv_var <- var(env_gv_cent)                          # Calculate the genetic covariance matrix

# --- SVD of the  variance-covariance matrix ---
svd_decomp <- svd(env_gv_var)
U_pop <- svd_decomp$u[, 1:k, drop = FALSE]         
D_pop <- diag(svd_decomp$d[1:k])                   

# Get latent covariates (H) and slopes (S)
H_pop <- U_pop %*% sqrt(D_pop)        # Covariates (PCA rotation)
S <- env_gv_cent %*% U_pop %*% solve(sqrt(D_pop))   # Slopes (PCA rotation)
range(S %*% t(H_pop) - env_gv_cent) # check for equivalence

# ---- Rotate the slopes and covariates around the TPE mean with drotate machinary ----
m <- colMeans(H_pop)
r1 <- as.numeric(m / sqrt(sum(m^2)))
P_perp <- diag(ncol(H_pop)) - r1 %*% t(r1)
ee <- eigen(P_perp, symmetric = TRUE)
R_rest <- ee$vectors[, ee$values > 1e-8, drop = FALSE]
R <- cbind(r1, R_rest)
H_rot <- H_pop %*% R
range(H_rot %*% t(H_rot) - env_gv_var) # check for equivalence

# now gather rotated covariates and slopes
D_dr <- diag(diag(t(H_rot) %*% H_rot))
H_dr <- H_rot %*% solve(sqrt(D_dr))  # new covariates
S_dr <- S %*% R %*% sqrt(D_dr)  # new slopes 
range(S_dr %*% t(H_dr) - env_gv_cent) # check for equivalence

# Organise into dataframe 
pop_df <- data.frame(id = DH@id, 
                    main_effect = rowMeans(env_gv), # Mean gv
                    main_effect_rot = S_dr[,1],    # First rotated slope (mean gv + all GEI variation perfectly associated with it)
                    pca_op = S[,1],                                                     # First slope based on a PCA 
                    var_raw = apply(env_gv_cent,1,var),                                 # Raw variance of gvs across TPE
                    var_scale = apply(scale(env_gv_cent),1,var),                        # Scaled variance of gvs across TPE
                    pca_rmsd = (sqrt(rowMeans((S[, -1] %*% t(H_pop[, -1]))^2))),        # RMSD based on the PCA rotation   
                    rot_rmsd = (sqrt(rowMeans((S_dr[, -1] %*% t(H_dr[, -1]))^2))),      # RMSD based on drotate machinery
                    rot_rmsd_kmax = (sqrt(rowMeans((S_dr[, 2:k_maxfit] %*% t(H_dr[, 2:k_maxfit]))^2))), # RMSD only considering first k_maxfit factors
                    DH@gv[,-1])

# Record genetic values in data frame 
df_pi = data.frame(DH@gv[,-1], strat = "popImp")




