/* 

  Linear regression in Picat.

  This program was inspired from the SetlX program regression.setlx

  This Picat model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Picat page: http://www.hakank.org/picat/

*/


main => go.

/*
n = 10
xs = [1,4,9,16,25,36,49,64,81,100]
ys = [1,8,27,64,125,216,343,512,729,1000]
-79.780220 + 9.929356*x (R: 0.987180)

n = 15
xs = [1,4,9,16,25,36,49,64,81,100,121,144,169,196,225]
ys = [1,8,27,64,125,216,343,512,729,1000,1331,1728,2197,2744,3375]
-249.160305 + 14.626939*x (R: 0.986688)

n = 30
xs = [1,4,9,16,25,36,49,64,81,100,121,144,169,196,225,256,289,324,361,400,441,484,529,576,625,676,729,784,841,900]
ys = [1,8,27,64,125,216,343,512,729,1000,1331,1728,2197,2744,3375,4096,4913,5832,6859,8000,9261,10648,12167,13824,15625,17576,19683,21952,24389,27000]
-1837.768924 + 28.699954*x (R: 0.986297)

*/
go =>  
  foreach (N in [10, 15, 30]) 
    printf("n = %d\n", N),
    Xs = [I**2 : I in 1 .. N],
    Ys = [I**3 : I in 1 .. N],
    printf("xs = %w\n", Xs),
    printf("ys = %w\n", Ys),
    [A,B,R] = regression(Xs,Ys),
    print_regression([A,B,R]),
    nl
  end,
  nl.

go2 =>
  N = 100000,
  _ = random2(),
  Ys = [I**2 : I in 1..N],
  print_regression(regression(Ys)),
  nl.


print_regression([A,B,R]) =>
   printf("%f + %f*X (R: %f)\n", A,B,R).


regression(Ys) = regression(1..Ys.len, Ys).

regression(Xs,Ys) = [A,B,R] =>
    A = regressionA(Xs, Ys),
    B = regressionB(Xs, Ys),
    R = correlation(Xs, Ys).


mean(Xs) = sum(Xs) / Xs.len.

variance(Xs) = Variance =>
    Mu = mean(Xs),
    N  = Xs.len,
    Variance = sum([ (X-Mu)**2 : X in Xs ]) / N.

squaredSum(Xs) = SS =>
    XMean = mean(Xs),
    SS = sum([ (X-XMean)**2 : X in Xs ]).

sumOfProds(Xs,Ys) = S =>
    XMean = mean(Xs),
    YMean = mean(Ys),
    S = sum([ (X-XMean)*(Y-YMean) : {X,Y} in zip(Xs,Ys)]).

regressionB(Xs,Ys) = sumOfProds(Xs, Ys) / squaredSum(Xs) + 0.0.

regressionA(Xs,Ys) = A =>
    XMean = mean(Xs),
    YMean = mean(Ys),
    A = YMean - regressionB(Xs, Ys) * XMean + 0.0.

% Correlations correficien (R)
correlation(Xs,Ys) = sumOfProds(Xs, Ys) / sqrt(squaredSum(Xs) * squaredSum(Ys)).