using LinearAlgebra
using DynamicPolynomials
using SwitchOnSafety
using Combinatorics
using SparseArrays
using JuMP, MosekTools
using SpecialFunctions




include("../src/RandomTrajectories.jl")
include("../src/AlgebraicLift.jl")
include("../src/ScenarioOpti.jl")
include("../src/ProbabilisticCertificates.jl")
include("../src/WhiteBoxAnalysis.jl")


dim = 6; numMode = 6; dimIn = Integer(floor(dim/2)+1)

numScen_budget = 40000


#beta = 0.01
#batchsize = 1000
tol=1e-3

A = [2*rand(Float64, (dim, dim)).-1 for i=1:numMode]#generate_switched_linear_systems(numMode,dim)
B = 2*rand(Float64, (dim, dimIn)).-1

#jsrboundopen= white_box_jsr(A)
#println("JSR open loop: $jsrboundopen")
gaWb = white_box_stabilization_quad(A,B)
println("White-box stabilization: $gaWb")

(state0_budget,state_budget) = generate_trajectories(1,A,numScen_budget)

jsrMat = []

K0 = zeros(dimIn,dim)

for N in 20000:1000:numScen_budget
    state0 = state0_budget[:,1:N]
    state = state_budget[:,1:N]
    K,jsr_bound = probabilistc_stability_certificate(state0,state;B=B,numMode=numMode,d=1,batchsize=1000,K0=K0,beta=0.01,tol=1e-3)
    push!(jsrMat,jsr_bound)
    println(jsrMat)
    println("White-box stabilization: $gaWb")
end




#n,M,m = 3,3,2
#Any[1.6905964015508155, 1.4949948318912094, 1.4578897728330178, 1.4449970570350583, 1.3661954310327227, 1.3557547083722759, 1.3519850357471235, 1.366442095956408, 1.3613413759324937, 1.3429185951062852, 1.3492377476476396, 1.353041725614804, 1.350881519910859, 1.351024882306086, 1.332142623005733, 1.3332099256037162, 1.3304895958561995, 1.3293168479495931, 1.3257398747945923, 1.3242527214300455]
#White-box stabilization: (1.3046225525634858, [-0.08704353863380242 -0.014885289542307709 -0.049119276805454554; -0.507914925287981 0.346359451257569 1.0595139493317263], [0.4730751424314884 0.05238239264562874 -0.11595043774240237; 0.05238239264562874 0.47706738397162957 -0.06627547776094828; -0.11595043774240237 -0.06627547776094828 0.5394594622004449])
#1000:1000:40000
#Any[1.5729841161676168, 1.0811554227453197, 1.1204878768144788, 0.9379157526294873, 1.0825047220475408, 0.9230141022655761, 0.9067207134210432, 0.9086742873941583, 0.9032261428457313, 0.9074603880590151, 0.9025114517834851, 0.9188062149754335, 0.879041002412628, 0.8870911981713733, 0.8926992150562685, 0.9244034116395815, 0.9213403526016543, 0.9186221154258161, 0.8857493686041322, 0.8845971086454935, 0.8796659173161556, 0.8787210471488426, 0.8778585397515366, 0.8770655056628548, 0.8763348706011853, 0.8756594770044023, 0.8816237449635299, 0.8810259381333286, 0.8803501376108269, 0.8799417718551741, 0.8794587848269098, 0.8790053814525175, 0.8785788973412633, 0.8752015696756477, 0.8748238216588682, 0.860271590366133, 0.8598788107531674, 0.8595062987770098, 0.8591372241094922, 0.8767243697570494]
#White-box stabilization: (0.8140390017107871, [0.4613815512089252 -0.5074407448126327 0.6781102279468121; -0.4054564358304964 0.9172867615097666 -0.5910686302058721], [0.1575553288978803 0.02705742403492229 -0.09788071638611223; 0.02705742403492229 0.3447930410671564 -0.04099837698782956; -0.09788071638611222 -0.04099837698782956 0.636873987762721])


#n,M,m = 4,4,3
#Any[Inf, Inf, Inf, Inf, 8.180216224894956, 6.799812206868395, 3.9527055267977347, 4.59174170956012, 3.0846800593837282, 2.6548561907838444, 2.511122991060089, 2.29880083863334, 2.161273525017097, 2.155982590641091, 2.0759785750943713, 2.6898719706804557, 2.3931499521503685, 2.048817168640575, 1.9826781075826072, 2.167094635668247]
#White-box stabilization: (1.2912508949641772, [-0.29467773970869054 0.45120789516804116 -0.044747113301878305 -1.3412127586682072; 0.29469719061888044 -0.12455867244325437 0.47017019594572057 -0.812536496343961; 0.46212127494925737 0.30630374660913356 0.03271426927287058 -1.5218113173193089], [0.3425385468310902 0.008731320963757413 0.13821082035445126 0.18629022929560726; 0.008731320963757404 0.27313481061590716 0.081610108315615 0.027369650943003167; 0.13821082035445126 0.08161010831561501 0.2951065970488265 0.08589689280299384; 0.1862902292956073 0.02736965094300317 0.08589689280299384 0.4106130096540772])        
#1000:1000:40000
#Any[Inf, Inf, 7.875913638698062, 7.17815675416856, 3.4416932736968016, 19.394289213014666, 2.583723965590616, 4.935647364390182, 1.984015772805858, 1.8674821390119283, 1.851433761685506, 2.2556567432257753, 2.1083617293205275, 2.0562587423170604, 1.720919007398438, 1.7461732618949821, 1.782734889401107, 1.7091518413970257, 1.7326683668613923, 1.6920558239323553, 1.7738463864124947, 1.8309626179003315, 1.7467378224996795, 1.6560635313126297, 1.650624547634913, 1.82414577393583, 1.6025444213063642, 1.6289801783733817, 1.745663331430371, 1.609830928758939, 1.597559164787974, 1.5904782457385636, 1.5838032484067968, 1.577498470309936, 1.5715323643953338, 1.6089367460744595, 1.553410839521955, 1.657015603207103, 1.5468122074359238, 1.546056871905595]
#White-box stabilization: (1.2337473357054138, [-0.007571043628938606 0.09068410059566147 -0.008151181227635546 -0.23763939318748917; 0.5431810653267845 0.25227603574433566 -0.1337727867212335 0.044797330378517915; 0.24885446041091436 0.4054576299301648 -0.3586486818168962 0.1656272016965966], [0.24836613958342035 0.09944045818289737 -0.16527107495145638 -0.12057885330835569; 0.09944045818289737 0.3576493794229119 -0.04223338581847068 -0.05938473714569832; -0.16527107495145635 -0.04223338581847068 0.428282720662584 0.07401563630157972; -0.12057885330835569 -0.05938473714569833 0.07401563630157973 0.5431548947140289])

#n,M,m = 4,6,3
#Any[Inf, Inf, Inf, 6.737095856110309, 3.675632937907526, 4.711663285904102, 3.6036786944899486, 3.245532901636809, 3.179705851172089, 2.928079175769009, 2.786632175891354, 2.5987061969012726, 2.690047270463758, 2.56824929071184, 2.3637115663198616, 3.0243999806253585, 3.0800995588355096, 2.7689719140449247, 2.640541451762192, 2.5793649540455488]
#White-box stabilization: (1.411688319459564, [0.18773192345444614 0.330514311888444 0.08303501025300539 0.6688093068290499; 0.096991858004839 -0.1376269451782156 0.16258760593673052 -0.28712484054435716; -0.05179889939849871 -0.21974552189934024 0.019847110276567503 -0.574729653178076], [0.3048746641659247 0.08704701508549992 -0.022825373701927137 0.03258573105931016; 0.08704701508549993 0.3155721631766691 -0.033540595048654746 0.01081474042260808; -0.022825373701927144 -0.033540595048654746 0.4273011813464707 -0.19692622371349458; 0.03258573105931017 0.01081474042260808 -0.19692622371349458 0.4285114388792512])
#1000:1000:40000
#Any[Inf, Inf, Inf, 25.084757321839195, 8.770136588631802, 9.976749384135678, 6.877930810422303, 6.939754176855923, 5.970890225057548, 5.158585131735316, 7.026520008542251, 3.9068211007996236, 2.7116440033548512, 2.6401441896367848, 2.6546224672270173, 2.479164798426021, 2.4708216864342907, 2.460497586092857, 2.429913835468321, 2.3758813963902914, 2.3461590747282215, 2.5836376981383578, 2.5429050265257245, 2.509091638025923, 2.4782728932067286, 2.4500540776504156, 2.4383169713775903, 2.4144773087877156, 2.3329332189333867, 2.251096846759749, 2.238871368744152, 2.2247113742266023, 2.17326016039637, 2.210486799215429, 2.198078155410338, 2.186357054727009, 2.2937024450686985, 2.268810905462152, 2.2561233423195572, 2.244082132489508]
#White-box stabilization: (1.8347012279895227, [1.2413485353963243 -0.13936635314223478 0.24391546443723508 -0.6272971813182449; 0.11882472083532031 0.006835856413333752 0.033442637097538024 -0.017462245766754272; 0.4572432136771731 -0.14943479419562547 0.11084913373271818 0.1330692144943799], [0.5288623013381373 -0.06955297816624266 0.07739492604379934 -0.14945226904731068; -0.06955297816624267 0.5133593931137286 0.024751588463856253 -0.06508656604630439; 0.07739492604379934 0.024751588463856253 0.4582612910506892 0.02558559352254509; -0.14945226904731068 -0.06508656604630439 0.02558559352254509 0.5703552027639498])


#n,M,m = 6,6,4
#30000:1000:40000
#ny[8.072715651100154, 7.768234690251989, 8.180570339667353, 7.886210946675976, 8.14582975086884, 7.871660888641335, 7.1033528127184, 7.390710250162713, 7.431287049154341, 5.567010043658964, 5.467431313638779]
#White-box stabilization: (1.8494425061800728, [0.012221018509647168 0.5312774075791682 0.038976706966226404 0.24115751812320538 0.29963266504137354 -0.24193177579705996; -0.28880253871014694 -0.23324503363629637 -0.030309376528313354 0.007505709936567113 -0.12073558590288352 0.15846373101244968; -0.20699525532892935 -0.4544479559088058 0.11982743997670932 -0.039258464791142526 -0.3181309400642283 0.3150143828761531; -0.0139831817182914 0.5874777519813043 -0.10178037874387551 0.12249922444250158 0.49381698904747273 -0.023525512008606104], [0.2629781490641416 0.0003384808471250383 -0.014510051589145008 -0.03191164850633556 0.013452938609376693 0.006781828774616212; 0.00033848084712503846 0.29904793232982807 0.0283758893731878 -0.00703246955168064 -0.05022634263568734 0.05898316898301458; -0.014510051589145012 0.0283758893731878 0.32192795235731525 -0.02850577360183669 -0.02556378959490166 0.02713605333508494; -0.03191164850633556 -0.007032469551680639 -0.028505773601836697 0.3370007449715669 -0.02596961925645344 0.06670614725088554; 0.013452938609376693 -0.05022634263568735 -0.02556378959490166 -0.02596961925645344 0.3207606635040778 -0.07275525717780275; 0.006781828774616212 0.05898316898301458 0.027136053335084936 0.06670614725088554 -0.07275525717780275 0.2882382320544729])
#20000:1000:40000
#Any[107.0843979270694, 59.342126663772156, 22.53861363197418, 19.35338764565481, 25.84527685861678, 21.968851042329128, 19.689300883307723, 49.614306286747734, 37.50510470585297, 10.732596162611243, 12.937874087207211, 8.87777827552819, 7.665576388216101, 7.2410821009579, 7.041489494504242, 6.996118755748454, 6.822953474887835, 6.505336929771839, 7.494983714450516, 6.310508993177824, 6.189122196716375]
#White-box stabilization: (1.9120856741024894, [0.2397554519049773 -0.10131199227907627 0.23889381969404355 -0.46808685505580094 -0.2830979655692059 -0.20420265706522728; 0.34396199478718237 0.04113041958533565 0.46428253981645234 -0.0810038935859582 -0.7453144451913242 -0.29377911940563817; 0.3294039638069217 -0.35268064974948693 0.423490743615612 -0.27429288114813877 -0.4997957061387201 -0.2536533942553235; 0.3459759820407083 -0.10603493805154796 0.38431492409875784 -0.16744635480413356 -1.045721513883651 0.0009893083224673024], [0.46126631090810905 0.0138872832302601 0.21125527828557872 0.12615429354805274 0.03787693562846611 -0.1251841989029774; 0.013887283230260098 0.41321690740308814 0.04124657326637017 -0.04856644977200873 0.05513722724728301 -0.11742322672944541; 0.21125527828557875 0.041246573266370164 0.6074761020842684 0.12138653021576508 -0.027014667936934723 -0.024310740534286227; 0.12615429354805274 -0.04856644977200873 0.12138653021576508 0.4420956139160551 0.06134324952849663 -0.07777124466288073; 0.03787693562846611 0.05513722724728302 -0.027014667936934723 0.061343249528496624 0.37137744074642076 -0.09637505917692687; -0.1251841989029774 -0.11742322672944543 -0.024310740534286227 -0.07777124466288071 -0.09637505917692687 0.44820687724218183])


