# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

EXPECTED_ANSWERS = [
    {  # test answers for PyTorch 1.6
        "integration_classification_2d": {
            "losses": [0.776835828070428, 0.1615355300011149, 0.07492854832938523, 0.04591309238865877],
            "best_metric": 0.9999184380485994,
            "infer_prop": [1029, 896, 980, 1033, 961, 1046],
        },
        "integration_segmentation_3d": {
            "losses": [
                0.5367561340332031,
                0.478084459900856,
                0.4581540793180466,
                0.44623913466930387,
                0.42341493666172025,
                0.42569945752620697,
            ],
            "best_metric": 0.9295084029436111,
            "infer_metric": 0.9296411260962486,
            "output_sums": [
                0.14302121377204619,
                0.15321686701244813,
                0.15267064069005093,
                0.1408481434833016,
                0.18862719991649474,
                0.16992848513054068,
                0.1479306037291329,
                0.1691071594535633,
                0.15804366588267224,
                0.18019304183940157,
                0.1635089455927468,
                0.16851606024285842,
                0.1454348651039073,
                0.11584957890961554,
                0.16255468027312903,
                0.20118089432240313,
                0.176187783307603,
                0.1004243279488101,
                0.19385348502657657,
                0.2030768555124136,
                0.196251372926592,
                0.20823046240222043,
                0.1631389353339986,
                0.13299661219478043,
                0.14917081129077908,
                0.14383374638201593,
                0.23050183928776746,
                0.1614747942341212,
                0.14913436515470202,
                0.10443081170610946,
                0.11978674347415241,
                0.13126176432899028,
                0.11570832453348577,
                0.15306806147195887,
                0.163673089782912,
                0.19394971756732426,
                0.22197501007172804,
                0.1812147930033603,
                0.19051659118682873,
                0.0774867922747158,
            ],
        },
        "integration_workflows": {
            "best_metric": 0.9219646483659745,
            "infer_metric": 0.921751058101654,
            "output_sums": [
                0.14183664321899414,
                0.1513957977294922,
                0.13804054260253906,
                0.13356828689575195,
                0.18456125259399414,
                0.16363763809204102,
                0.14090299606323242,
                0.16649389266967773,
                0.15651893615722656,
                0.17655134201049805,
                0.16116666793823242,
                0.1644763946533203,
                0.14383649826049805,
                0.11055326461791992,
                0.16080379486083984,
                0.19629907608032227,
                0.17441415786743164,
                0.053577423095703125,
                0.19043779373168945,
                0.19904804229736328,
                0.19526052474975586,
                0.20304107666015625,
                0.16030025482177734,
                0.13170623779296875,
                0.15118932723999023,
                0.13686418533325195,
                0.22668886184692383,
                0.1611471176147461,
                0.1472463607788086,
                0.10427379608154297,
                0.11962461471557617,
                0.1305704116821289,
                0.11204910278320312,
                0.15171337127685547,
                0.15962505340576172,
                0.18976259231567383,
                0.21649456024169922,
                0.17761802673339844,
                0.18516874313354492,
                0.03636503219604492,
            ],
            "best_metric_2": 0.9219559609889985,
            "infer_metric_2": 0.9217371672391892,
            "output_sums_2": [
                0.14187288284301758,
                0.15140819549560547,
                0.13802719116210938,
                0.1335887908935547,
                0.18454980850219727,
                0.1636652946472168,
                0.14091157913208008,
                0.16653108596801758,
                0.15651702880859375,
                0.17658615112304688,
                0.1611957550048828,
                0.16448307037353516,
                0.14385128021240234,
                0.1105203628540039,
                0.16085100173950195,
                0.19626951217651367,
                0.17442035675048828,
                0.053586483001708984,
                0.19042730331420898,
                0.1990523338317871,
                0.1952815055847168,
                0.20303773880004883,
                0.16034317016601562,
                0.13172531127929688,
                0.15118741989135742,
                0.1368694305419922,
                0.22667837142944336,
                0.16119050979614258,
                0.14726591110229492,
                0.10426473617553711,
                0.11961841583251953,
                0.13054800033569336,
                0.11203193664550781,
                0.15172529220581055,
                0.15963029861450195,
                0.18975019454956055,
                0.21646499633789062,
                0.17763566970825195,
                0.18517112731933594,
                0.03638744354248047,
            ],
        },
    },
    {  # test answers for PyTorch 1.7
        "integration_classification_2d": {
            "losses": [0.777176220515731, 0.16019743723664315, 0.07480076164197011, 0.045643698364780966],
            "best_metric": 0.9999418774120775,
            "infer_prop": [1030, 897, 980, 1033, 960, 1048],
        },
        "integration_segmentation_3d": {
            "losses": [
                0.5427072256803512,
                0.46434969305992124,
                0.45358552038669586,
                0.4363856494426727,
                0.42080804109573366,
                0.42058534920215607,
            ],
            "best_metric": 0.9292903542518616,
            "infer_metric": 0.9306288316845894,
            "output_sums": [
                0.14192493409895743,
                0.15182314591386872,
                0.15143080738742032,
                0.13972497034181824,
                0.18790884439406313,
                0.16933812661492562,
                0.14664343345928132,
                0.1678599094806423,
                0.1568852615222309,
                0.17882538307200632,
                0.16226220644853354,
                0.16756325103417588,
                0.1449974856885373,
                0.1160602083671129,
                0.1614830941632057,
                0.20060717335382267,
                0.17543495742507476,
                0.10308107883493946,
                0.19289222718691168,
                0.20225689438356148,
                0.19587806881756237,
                0.20773073456322155,
                0.16193015294299506,
                0.13181961683097554,
                0.14850995284454005,
                0.14238637655756,
                0.2307113922277095,
                0.1608335768948913,
                0.1480752874532259,
                0.1038477413165911,
                0.11880665574424197,
                0.13084873656303445,
                0.1141965805147642,
                0.1531586543003841,
                0.16275008603701097,
                0.19320476187766733,
                0.2217811250932611,
                0.18027048819200148,
                0.18958803602663193,
                0.08653716931250294,
            ],
        },
        "integration_workflows": {
            "best_metric": 0.9217087924480438,
            "infer_metric": 0.9214379042387009,
            "output_sums": [
                0.14209461212158203,
                0.15126705169677734,
                0.13800382614135742,
                0.1338181495666504,
                0.1850571632385254,
                0.16372442245483398,
                0.14059066772460938,
                0.16674423217773438,
                0.15653657913208008,
                0.17690563201904297,
                0.16154909133911133,
                0.16521310806274414,
                0.14388608932495117,
                0.1103353500366211,
                0.1609959602355957,
                0.1967010498046875,
                0.1746964454650879,
                0.05329275131225586,
                0.19098854064941406,
                0.19976520538330078,
                0.19576644897460938,
                0.20346736907958984,
                0.1601848602294922,
                0.1316051483154297,
                0.1511220932006836,
                0.13670969009399414,
                0.2276287078857422,
                0.1611800193786621,
                0.14751672744750977,
                0.10413789749145508,
                0.11944007873535156,
                0.1305546760559082,
                0.11204719543457031,
                0.15145111083984375,
                0.16007614135742188,
                0.1904129981994629,
                0.21741962432861328,
                0.17812013626098633,
                0.18587207794189453,
                0.03605222702026367,
            ],
            "best_metric_2": 0.9210659921169281,
            "infer_metric_2": 0.9208109736442566,
            "output_sums_2": [
                0.14227628707885742,
                0.1515035629272461,
                0.13819408416748047,
                0.13402271270751953,
                0.18525266647338867,
                0.16388607025146484,
                0.14076614379882812,
                0.16694307327270508,
                0.15677356719970703,
                0.1771831512451172,
                0.16172313690185547,
                0.1653728485107422,
                0.14413118362426758,
                0.11057281494140625,
                0.16121912002563477,
                0.19680166244506836,
                0.1748638153076172,
                0.053426265716552734,
                0.19117307662963867,
                0.19996356964111328,
                0.1959366798400879,
                0.20363712310791016,
                0.16037797927856445,
                0.13180780410766602,
                0.1513657569885254,
                0.13686084747314453,
                0.2277364730834961,
                0.16137409210205078,
                0.1476879119873047,
                0.10438394546508789,
                0.11967992782592773,
                0.13080739974975586,
                0.11226606369018555,
                0.15168476104736328,
                0.1602616310119629,
                0.190582275390625,
                0.21756458282470703,
                0.17825984954833984,
                0.18604803085327148,
                0.036206722259521484,
            ],
        },
    },
    {  # test answers for PyTorch 21.04, cuda 11.3
        "integration_classification_2d": {
            "losses": [0.7772567988770782, 0.16357883198815545, 0.0748426011840629, 0.045560025710873545],
            "best_metric": 0.9999362036681547,
            "infer_prop": [1030, 898, 981, 1033, 960, 1046],
        },
        "integration_segmentation_3d": {
            "losses": [
                0.5462346076965332,
                0.4699550330638885,
                0.4407052755355835,
                0.4473582059144974,
                0.4345871120691299,
                0.4268435090780258,
            ],
            "best_metric": 0.9325245052576066,
            "infer_metric": 0.9326683700084686,
            "output_sums": [
                0.14224469870198278,
                0.15221021012369151,
                0.15124158255724182,
                0.13988812880932433,
                0.18869885039284465,
                0.16944664085835437,
                0.14679946398855015,
                0.1681337815374021,
                0.1572538225010156,
                0.179386563044054,
                0.162734465243387,
                0.16831902111202945,
                0.1447043535420074,
                0.11343210557896033,
                0.16199135405262954,
                0.20095180481987404,
                0.17613484080473857,
                0.09717457016552708,
                0.1940439758638305,
                0.2033698355271389,
                0.19628583555443793,
                0.20852096425983455,
                0.16202004771083997,
                0.13206408917949392,
                0.14840973098125526,
                0.14237425379050472,
                0.23165483128059614,
                0.16098621485325398,
                0.14831028015056963,
                0.10317099380415945,
                0.118716576251689,
                0.13002315213569166,
                0.11436407827087304,
                0.1522274707636008,
                0.16314910792851098,
                0.1941135852761834,
                0.22309890968242424,
                0.18111804948625987,
                0.19043976068601465,
                0.07442812452084423,
            ],
        },
    },
    {  # test answers for PyTorch 1.9
        "integration_workflows": {
            "output_sums_2": [
                0.14213180541992188,
                0.15153264999389648,
                0.13801145553588867,
                0.1338348388671875,
                0.18515968322753906,
                0.16404008865356445,
                0.14110612869262695,
                0.16686391830444336,
                0.15673542022705078,
                0.1772594451904297,
                0.16174745559692383,
                0.16518878936767578,
                0.1440296173095703,
                0.11033201217651367,
                0.1611781120300293,
                0.19660568237304688,
                0.17468547821044922,
                0.053053855895996094,
                0.1909656524658203,
                0.19952869415283203,
                0.1957845687866211,
                0.2034916877746582,
                0.16042661666870117,
                0.13193607330322266,
                0.15104389190673828,
                0.13695430755615234,
                0.22720861434936523,
                0.16157913208007812,
                0.14759159088134766,
                0.10379791259765625,
                0.11937189102172852,
                0.1306462287902832,
                0.11205482482910156,
                0.15182113647460938,
                0.16006708145141602,
                0.19011592864990234,
                0.21713829040527344,
                0.17794132232666016,
                0.18584394454956055,
                0.03577899932861328,
            ]
        },
        "integration_segmentation_3d": {  # for the mixed readers
            "losses": [
                0.5645154356956482,
                0.4984356611967087,
                0.472334086894989,
                0.47419720590114595,
                0.45881829261779783,
                0.43097741305828097,
            ],
            "best_metric": 0.9325698614120483,
            "infer_metric": 0.9326590299606323,
        },
    },
    {  # test answers for PyTorch 21.10
        "integration_classification_2d": {
            "losses": [0.7806222991199251, 0.16259610306495315, 0.07529311385124353, 0.04640352608529246],
            "best_metric": 0.9999369155431564,
            "infer_prop": [1030, 898, 981, 1033, 960, 1046],
        },
        "integration_segmentation_3d": {
            "losses": [
                0.5462362408638001,
                0.4913381844758987,
                0.4526856362819672,
                0.43404580652713776,
                0.42532919645309447,
                0.4160102754831314,
            ],
            "best_metric": 0.9357608556747437,
            "infer_metric": 0.9359462857246399,
            "output_sums": [
                0.14133183650702907,
                0.15129517085134564,
                0.15039408698301698,
                0.1388800895551786,
                0.18765019147239637,
                0.16847158867677473,
                0.14567945622102715,
                0.16728557092807228,
                0.15601444057659314,
                0.17816339678760573,
                0.1616256801482474,
                0.16733042976922818,
                0.14342795433701588,
                0.1122946416901734,
                0.16105778942392063,
                0.20017543167070598,
                0.17512204704647916,
                0.09592956823274325,
                0.19316383411238341,
                0.2022308530579937,
                0.19527218778022315,
                0.2075871950564991,
                0.16083565516485876,
                0.13111518931029637,
                0.1473909261474288,
                0.14161210629657228,
                0.23102446985179093,
                0.15980667305916593,
                0.14760356792082058,
                0.1018092235719272,
                0.11792260857122504,
                0.1285278390386459,
                0.11275165891441473,
                0.15101653432548032,
                0.16236351926994622,
                0.1932631773335222,
                0.2221395787381994,
                0.18003549292918666,
                0.18940543270178078,
                0.07430261166443994,
            ],
        },
        "integration_workflows": {
            "output_sums": [
                0.14211511611938477,
                0.1516571044921875,
                0.1381092071533203,
                0.13403034210205078,
                0.18480682373046875,
                0.16382598876953125,
                0.14140796661376953,
                0.1665945053100586,
                0.15700864791870117,
                0.17697620391845703,
                0.16163396835327148,
                0.16488313674926758,
                0.1442713737487793,
                0.11060476303100586,
                0.16111087799072266,
                0.19617986679077148,
                0.1744403839111328,
                0.052786827087402344,
                0.19046974182128906,
                0.19913578033447266,
                0.19527721405029297,
                0.2032318115234375,
                0.16050148010253906,
                0.13228464126586914,
                0.1512293815612793,
                0.1372208595275879,
                0.22692251205444336,
                0.16164922714233398,
                0.14729642868041992,
                0.10398292541503906,
                0.1195836067199707,
                0.13096046447753906,
                0.11221647262573242,
                0.1521167755126953,
                0.1599421501159668,
                0.1898345947265625,
                0.21675777435302734,
                0.1777491569519043,
                0.18526840209960938,
                0.035144805908203125,
            ],
            "output_sums_2": [
                0.14200592041015625,
                0.15146303176879883,
                0.13796186447143555,
                0.1339101791381836,
                0.18489742279052734,
                0.1637406349182129,
                0.14113903045654297,
                0.16657161712646484,
                0.15676355361938477,
                0.17683839797973633,
                0.1614980697631836,
                0.16493558883666992,
                0.14408016204833984,
                0.11035394668579102,
                0.1610560417175293,
                0.1962742805480957,
                0.17439842224121094,
                0.05285835266113281,
                0.19057941436767578,
                0.19914865493774414,
                0.19533538818359375,
                0.20333576202392578,
                0.16032838821411133,
                0.13197898864746094,
                0.1510462760925293,
                0.13703680038452148,
                0.2270984649658203,
                0.16144943237304688,
                0.1472611427307129,
                0.10393238067626953,
                0.11940813064575195,
                0.1307811737060547,
                0.11203241348266602,
                0.15186500549316406,
                0.15992307662963867,
                0.18991422653198242,
                0.21689796447753906,
                0.1777033805847168,
                0.18547868728637695,
                0.035192012786865234,
            ],
        },
    },
]


def test_integration_value(test_name, key, data, rtol=1e-2):
    for (idx, expected) in enumerate(EXPECTED_ANSWERS):
        if test_name not in expected:
            continue
        if key not in expected[test_name]:
            continue
        value = expected[test_name][key]
        if np.allclose(data, value, rtol=rtol):
            print(f"matched {idx} result of {test_name}, {key}, {rtol}.")
            return True
    raise ValueError(f"no matched results for {test_name}, {key}. {data}.")
