using ACE1
using IPFitting

al = IPFitting.Data.read_xyz(@__DIR__() * "/TiAl_tutorial_DB.xyz", energy_key="energy", force_key="force", virial_key="virial")[1:10:end];

r0 = 0.5*(rnn(:Ti) + rnn(:Al))

ACE_B = ACE1.Utils.rpi_basis(species = [:Ti, :Al],
                              N = 3,
                              r0 = r0,
                              rin = 0.6 * r0,
                              rcut = 5.5,
                              maxdeg = 6)
                              
Bpair = pair_basis(species = [:Ti, :Al],
      r0 = r0,
      maxdeg = 6,
      rcut = 7.0,
      pcut = 1,
      pin = 0)  

B = JuLIP.MLIPs.IPSuperBasis([Bpair, ACE_B]);

dB = LsqDB("", B, al)

Vref = OneBody(:Ti => -1586.0195, :Al => -105.5954)

weights = Dict(
        "FLD_TiAl" => Dict("E" => 30.0, "F" => 1.0 , "V" => 1.0 ),
        "TiAl_T5000" => Dict("E" => 5.0, "F" => 1.0 , "V" => 1.0 ))

solver_type = :ard
laplace_precon = true

if solver_type == :lsqr
	solver = Dict(
        	"solver" => :lsqr,
        	"lsqr_damp" => 1e-2,
        	"lsqr_atol" => 1e-6)
elseif solver_type == :rrqr
	solver = Dict(
        	"solver" => :rrqr,
        	"rrqr_tol" => 1e-5)
elseif solver_type == :brr
	solver = Dict(
        	"solver" => :brr,
		"brr_tol" => 1e-3)
elseif solver_type == :ard 
	solver= Dict(
         	"solver" => :ard,
         	"ard_tol" => 1e-3,
         	"ard_threshold_lambda" => 10000)
end

if laplace_precon
	using LinearAlgebra
	rlap_scal = 3.0
	P = Diagonal(vcat(ACE1.scaling.(dB.basis.BB, rlap_scal)...))
	solver["P"] = P
end

IP, lsqinfo = lsqfit(dB, solver=solver, weights=weights, Vref=Vref, error_table = true);

rmse_table(lsqinfo["errors"])

save_dict("./TiAl_tutorial_pot.json", Dict("IP" => write_dict(IP), "info" => lsqinfo))

ACE1.ExportMulti.export_ACE("./TiAl_tutorial_pot.yace", IP)
