# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.


from io import StringIO
import os
import csv
from typing import List, Any

from .data_types import CO3DTask, CO3DSequenceSet


BLANK_PREDICTION_RESULTS = {}


def _read_result_csv(s: str):
    # with open(fl, "r") as f:
    f = StringIO(s)
    csvreader = csv.reader(f)
    rows = [row for row in csvreader]
    rows = rows[1:]
    header = rows[0]
    data = rows[1:-1]
    def _getcol(col_name: str, row: List[Any]) -> Any:
        c = row[header.index(col_name)]
        try:
            return float(c)
        except:
            return c
    parsed = {
        (_getcol("Category", r), _getcol("Subset name", r)): {
            k: _getcol(k, r) for k in header
        } for r in data
    }
    return parsed


CSVs = {
"fewview_dev": """
Category,Subset name,psnr_masked,psnr_fg,psnr_full_image,depth_abs_fg,iou
apple,fewview_dev,18.40938866633708,6.884780900276403,5.732459292886711,0.48950375965076004,0.0
backpack,fewview_dev,18.375179837644755,11.884768822089297,5.492699127831022,0.580444590643848,0.0
ball,fewview_dev,15.65596825167019,5.697924649467918,5.391241119316918,0.43991856992712286,0.0
banana,fewview_dev,18.461971791362227,6.118058719441003,5.8697287026999625,0.5312080518960041,0.0
baseballbat,fewview_dev,20.451565072548348,6.7702838462526325,6.133595679990426,0.787964382936369,0.0
baseballglove,fewview_dev,15.899123723379235,8.491206359449485,5.952075366998026,0.5044438522210485,0.0
bench,fewview_dev,13.835660454286623,6.1021708060060185,5.338972434739994,0.8728473659927769,0.0
bicycle,fewview_dev,14.85079899106894,7.178515383648441,5.4468849723020165,0.7596495667817377,0.0
book,fewview_dev,13.526778301589218,5.929520397898452,6.1038428839075625,0.7119168557685552,0.0
bottle,fewview_dev,17.756936987543572,7.695879675777415,5.792669536453962,1.1126274259151023,0.0
bowl,fewview_dev,12.117324340446702,3.522034136500667,6.132690804727037,0.548212652825193,0.0
broccoli,fewview_dev,17.60270342882336,8.135587140185267,5.636059385848195,0.48109570750702163,0.0
cake,fewview_dev,14.831394456777907,6.641730746137352,5.778288244687103,0.4713467452914664,0.0
car,fewview_dev,12.199833440326447,6.2695458065545955,5.7147062915561,0.6731242096715442,0.0
carrot,fewview_dev,18.42032666772822,6.336027619876071,5.2655157144357,0.7425826445279987,0.0
cellphone,fewview_dev,18.54815997270957,9.132949039155196,5.920507132031587,0.7256476083461838,0.0
chair,fewview_dev,14.254104990224922,6.8885175096457525,5.42230365019509,0.8701949198272996,0.0
couch,fewview_dev,12.096141908081652,8.498063614467037,6.839693292778098,0.6672055849897333,0.0
cup,fewview_dev,16.30300593190912,6.263725950094426,5.419278138684526,1.109737605178693,0.0
donut,fewview_dev,17.760249549810045,7.19401090262162,5.406775287613137,0.5831024075924244,0.0
frisbee,fewview_dev,13.006974807290442,5.348851057119092,6.081314892526941,0.6282357528069842,0.0
hairdryer,fewview_dev,18.307693246477385,7.653327373043194,5.796698293526376,0.5692578716769887,0.0
handbag,fewview_dev,16.863888776603684,9.668777191048893,5.885582988575421,0.6140305534695657,0.0
hotdog,fewview_dev,16.576000201780598,6.7813353163227275,6.479828364566311,0.5515738226619902,0.0
hydrant,fewview_dev,14.35863704229326,5.557106534568748,5.486735221682155,0.7370800150837736,0.0
keyboard,fewview_dev,18.319239151881423,10.9398173290579,5.471888028766401,0.591969625411462,0.0
kite,fewview_dev,13.759580600059902,6.095096560743659,5.5797533716568335,0.3686704352187232,0.0
laptop,fewview_dev,17.958107529829775,10.58932076091378,5.9870485037655365,0.6760399403943799,0.0
microwave,fewview_dev,12.641232654595555,7.5579894876019935,5.7736075695959785,0.7816656712123962,0.0
motorcycle,fewview_dev,13.902730964332383,7.837737363341203,5.6993349939287,0.8026270041676278,0.0
mouse,fewview_dev,22.139654039699753,11.380540045528843,5.26534717648027,0.6258851366555073,0.0
orange,fewview_dev,16.965398815565717,5.392140191707388,5.868309801114943,0.45518186645635506,0.0
parkingmeter,fewview_dev,17.969579417828633,8.303453741571293,5.550653705252322,2.7703986799279625,0.0
pizza,fewview_dev,14.044388259713267,6.467125499434811,6.349638051827558,0.5445261030741094,0.0
plant,fewview_dev,15.912698636112678,8.209728015160032,5.41847542705161,0.9729385734872266,0.0
remote,fewview_dev,18.901389746835065,8.809855001539868,5.6508358729724995,0.5809070430213752,0.0
sandwich,fewview_dev,14.961081916655587,5.359419050654777,6.486182655727676,0.5273259918746086,0.0
skateboard,fewview_dev,15.12940600031295,6.633805444460857,6.075841409914119,0.5708218125938797,0.0
stopsign,fewview_dev,18.52676122564753,6.61671306856769,5.412139613407474,6.290707304470178,0.0
suitcase,fewview_dev,16.493029339685542,10.757954804495968,6.232275999259873,0.5967537541074001,0.0
teddybear,fewview_dev,12.49373038673622,5.562061567728542,5.8834174182726855,0.6012993745910462,0.0
toaster,fewview_dev,15.590308176317933,8.571510283192422,5.8223530170835565,0.7087675899756055,0.0
toilet,fewview_dev,11.053325723237059,3.745954412389449,5.831752233322646,0.7324808735388084,0.0
toybus,fewview_dev,15.74397288343334,5.87386919966778,5.694742423634763,0.644572040998336,0.0
toyplane,fewview_dev,15.271423476084475,4.920347774565625,5.603913746917713,0.5686183372913356,0.0
toytrain,fewview_dev,19.250492955217194,8.365187557837626,5.5957012947860445,0.6429103676877059,0.0
toytruck,fewview_dev,15.813126824200825,7.367196186168707,5.59287438907558,0.5748745851615271,0.0
tv,fewview_dev,18.455985344741848,11.821412211957313,5.87636504861574,0.6193668766022515,0.0
umbrella,fewview_dev,13.388214509185625,6.669691460242465,5.398996667950242,0.5547154568934756,0.0
vase,fewview_dev,17.385895374160103,7.695607020715037,5.667400967410725,1.0544596567185702,0.0
wineglass,fewview_dev,14.92593215613611,5.489494483032894,5.883318241506832,2.09036588666451,0.0
MEAN,-,16.028754842096472,7.3270142749005025,5.768476753918801,0.8374863237526772,0.0
""",
"fewview_test": """
Category,Subset name,psnr_masked,psnr_fg,psnr_full_image,depth_abs_fg,iou
apple,fewview_test,18.51983235506069,6.710896207691665,5.622396257710374,0.45868530307683764,0.0
backpack,fewview_test,15.329616295156082,9.704246779430184,6.021398266902823,0.5274631579925675,0.0
ball,fewview_test,16.999140797902346,6.393148333684946,6.167099298585788,0.42074640466733093,0.0
banana,fewview_test,17.20449002482513,6.2347690237546765,5.337301584435467,0.5906480660508107,0.0
baseballbat,fewview_test,20.598735999896142,6.724621984421882,5.929346230877072,0.46383516633969724,0.0
baseballglove,fewview_test,16.250018316676424,8.485414452103313,5.35050821728197,0.5755057054113818,0.0
bench,fewview_test,13.380691505741307,6.217615311139159,5.389882231932645,0.8591881917970259,0.0
bicycle,fewview_test,15.804150486121728,8.539006404409536,7.293404052140095,0.7740858337090635,0.0
book,fewview_test,14.350489743207989,5.356299926470255,5.138131270946916,0.6249600811612394,0.0
bottle,fewview_test,17.257503711230473,7.332068784914889,5.825424785199224,1.0062512850600411,0.0
bowl,fewview_test,12.7586871865527,5.952472495887487,7.350451995400975,0.7734948803009338,0.0
broccoli,fewview_test,17.69069033947863,8.250871950138103,5.718669980890903,0.5437043438960382,0.0
cake,fewview_test,14.809462963950144,6.142164342026519,6.145654847812541,0.45489466623242036,0.0
car,fewview_test,11.914391205648087,6.5335541836879925,5.90360267479956,0.9021454444786102,0.0
carrot,fewview_test,20.060924545297425,6.219697054467009,5.261149123525815,0.7081597814658059,0.0
cellphone,fewview_test,21.520117285013956,10.847631110964242,5.41747877060995,1.0517241006106035,0.0
chair,fewview_test,14.691657730804202,8.959579180137167,6.878377818012938,0.8045192519054911,0.0
couch,fewview_test,11.545670382508696,8.419983656626247,6.902446179473004,0.6761085327114593,0.0
cup,fewview_test,17.79448614165711,6.495705819546957,5.5050360165654855,0.8834131631626546,0.0
donut,fewview_test,18.596152225400257,6.892531195772306,6.240000810567556,0.5443665622620474,0.0
frisbee,fewview_test,14.370690470903668,6.048295011020775,6.136056575421687,0.4830201400666513,0.0
hairdryer,fewview_test,18.47390481689051,7.494774772300304,5.743646634555602,0.5239972887128962,0.0
handbag,fewview_test,13.87987101022776,8.280409779606966,6.572322491579377,0.6866448922525301,0.0
hotdog,fewview_test,18.436410464732152,7.713564800659037,5.859372904290447,0.5873852722036716,0.0
hydrant,fewview_test,14.768617799865435,5.67036284794227,5.71565321761019,0.9328092564314482,0.0
keyboard,fewview_test,18.875163364703024,10.97846088231997,5.392007807994692,0.42114457863505195,0.0
kite,fewview_test,12.882975207164943,6.079375329369365,5.243720977367847,0.571440938913041,0.0
laptop,fewview_test,16.68965246676936,9.765618650745138,6.127183977142236,0.8968296529628422,0.0
microwave,fewview_test,13.859058432153368,8.649172226048128,6.809269971869398,0.8740670698190732,0.0
motorcycle,fewview_test,12.922201328542098,7.659321482648036,5.3469570020173816,0.7923491167407205,0.0
mouse,fewview_test,25.03083236821661,10.870194079196883,5.61381320415904,0.5803283306516662,0.0
orange,fewview_test,17.906264108511905,5.863058031859002,5.902648030774557,0.4927651700044394,0.0
parkingmeter,fewview_test,24.486359595107576,10.777998512312754,4.875545759481984,3.9189161735406275,0.0
pizza,fewview_test,15.25053153218815,6.195657831341678,5.888809317232928,0.5366542850357786,0.0
plant,fewview_test,14.533347345876026,8.213483475587314,5.9657101837783895,0.8745105580745663,0.0
remote,fewview_test,18.685696193857062,9.167126712684974,5.283444994288521,0.5784209284648094,0.0
sandwich,fewview_test,14.954638830523134,5.489779040424508,6.203690658497073,0.582476274688696,0.0
skateboard,fewview_test,18.921604245076754,8.111335322871586,4.540996792864179,0.8144729054641098,0.0
stopsign,fewview_test,20.83021952727707,7.7066182145576425,5.596606825038416,6.195708155269956,0.0
suitcase,fewview_test,14.568523293458965,8.872585021337093,5.526936386940414,0.5437482494754128,0.0
teddybear,fewview_test,13.184137897313038,5.667378086474551,5.638538121962938,0.6289599526865502,0.0
toaster,fewview_test,15.398766247640951,8.138341096517484,6.073562974743127,0.7335666912630792,0.0
toilet,fewview_test,10.138714105703048,3.8756171226863025,5.85450160774978,0.7892172212095283,0.0
toybus,fewview_test,15.925097991923954,6.517829456639026,5.691133527297476,0.6022958688384993,0.0
toyplane,fewview_test,16.703705769834098,5.323541429433026,5.46165954412417,0.5639341931778066,0.0
toytrain,fewview_test,17.859279914562713,7.8933999002371715,5.604032948369101,0.6932112812874591,0.0
toytruck,fewview_test,16.971557700694344,7.745719186191729,5.794916102483104,0.564653671235697,0.0
tv,fewview_test,18.037750946556894,13.741247943038163,8.747561838523023,0.5162819237405952,0.0
umbrella,fewview_test,13.092407842058238,6.756963662911218,5.447907114523638,0.534506784839016,0.0
vase,fewview_test,18.54297573271471,8.090029952142554,5.668374190385807,0.84122947818443,0.0
wineglass,fewview_test,16.386668940524114,5.5524702294978345,5.735686759902533,1.4353355366647544,0.0
MEAN,-,16.463618328111792,7.555333495840728,5.871765271698825,0.8516623875064206,0.0
""",
"manyview_dev": """
Category,Subset name,psnr_masked,psnr_fg,psnr_full_image,depth_abs_fg,iou
apple,manyview_dev_0,18.264030492114536,8.350223131127144,4.366539721003419,0.4195637484678012,0.0
apple,manyview_dev_1,14.137138507072345,6.6045994842301345,6.240087240624211,0.43567804409070654,0.0
ball,manyview_dev_0,14.673712693605873,6.091306495279248,5.217217027846326,0.35927968102112323,0.0
ball,manyview_dev_1,11.090845071075146,4.64095367064294,2.463653189968876,0.30228020972164427,0.0
bench,manyview_dev_0,13.333540945296608,4.137188797564715,3.844656341335867,0.8008696769825814,0.0
bench,manyview_dev_1,11.474174975542255,3.892151505117967,4.14563643434561,0.8577265682977291,0.0
book,manyview_dev_0,13.964168705937992,5.302433873449493,5.950633752149304,0.668803861808978,0.0
book,manyview_dev_1,12.398406799192342,4.119572830245314,6.039375672561894,0.8608240982086351,0.0
bowl,manyview_dev_0,16.958798002755774,4.9461020198227335,5.578702964374623,0.6690737351712432,0.0
bowl,manyview_dev_1,12.420483353954074,5.756645234213993,6.069489156010504,0.5819949787763078,0.0
broccoli,manyview_dev_0,19.630737300870244,9.406282525085935,6.402535226376115,0.7907156923061898,0.0
broccoli,manyview_dev_1,18.781287064441447,8.09672300742875,4.67134680549106,0.4626196557341922,0.0
cake,manyview_dev_0,14.799043006158593,5.867235047104056,5.7329760554862945,0.5205964759006821,0.0
cake,manyview_dev_1,17.84162321617,9.41822453353167,3.7158681607815254,0.3612821873000541,0.0
donut,manyview_dev_0,19.315033141413654,9.455566547834058,3.910254156226572,0.5413953368124613,0.0
donut,manyview_dev_1,22.26734997183049,10.174649831308487,4.199195894665875,0.5521516658527057,0.0
hydrant,manyview_dev_0,14.599159376924849,5.655154414726878,5.289620369144585,0.9737327772204973,0.0
hydrant,manyview_dev_1,14.544431000855953,5.876377992594626,4.506377178812374,1.0210153410111495,0.0
mouse,manyview_dev_0,22.553107676356586,12.793445604091437,5.927286492328659,0.5816200334131308,0.0
mouse,manyview_dev_1,17.89414321396086,8.956320087603723,7.097351162295129,0.5222896946353802,0.0
orange,manyview_dev_0,13.732343455171254,5.052956697685929,5.679024711561304,0.40213060027513875,0.0
orange,manyview_dev_1,14.71190574360874,4.956667990371484,5.836996460679712,0.43328379232231895,0.0
plant,manyview_dev_0,17.56722473025224,10.851111767732277,6.940102616941581,0.9601928359930311,0.0
plant,manyview_dev_1,18.62091024389777,11.114146143571679,8.919832772445316,0.845715675126882,0.0
remote,manyview_dev_0,12.004470911615606,2.3372367853347664,5.928692360063941,0.6355222400483482,0.0
remote,manyview_dev_1,13.035720177392095,4.368321832863184,3.7645273565115303,0.6257342864206513,0.0
skateboard,manyview_dev_0,14.087374862144243,6.183930758291541,7.7026533167035085,0.7381270587952287,0.0
skateboard,manyview_dev_1,15.24606555170737,6.935641480347134,6.728247832458047,0.6846367731825937,0.0
suitcase,manyview_dev_0,13.819257223346327,5.727869083939035,5.9663188950446795,0.42728104332046707,0.0
suitcase,manyview_dev_1,23.33527836247522,12.70130752964975,5.440617175698944,0.7376517524662343,0.0
teddybear,manyview_dev_0,15.310590723595963,7.5183318102880765,5.187722505560557,0.6132311702409632,0.0
teddybear,manyview_dev_1,19.00287693135702,11.380410989980264,5.372428296399181,0.655451568067443,0.0
toaster,manyview_dev_0,16.09490094737935,7.357336873218335,5.733018822009381,0.6335824697011363,0.0
toaster,manyview_dev_1,13.391233953784758,6.32606222531527,6.035255066975607,0.7543408733149064,0.0
toytrain,manyview_dev_0,14.60365232137707,8.252354438191217,7.28055045581793,0.5177963318470418,0.0
toytrain,manyview_dev_1,20.508004149463403,10.310151926704073,8.745624247957407,0.4164560185628414,0.0
toytruck,manyview_dev_0,18.495843812347488,9.077851138541167,4.742593752879244,0.8234759152694971,0.0
toytruck,manyview_dev_1,12.550467820571148,5.368998580430165,6.689171662380995,0.581289871598415,0.0
vase,manyview_dev_0,18.188943183563104,9.441252383753767,3.3505357321672142,0.7542355580664746,0.0
vase,manyview_dev_1,18.434184156563,9.303826519080554,6.071437833814365,0.9019223769623579,0.0
MEAN,-,16.092061594428568,7.352673089707325,5.58710387189748,0.635639291857879,0.0
""",
"manyview_test": """
Category,Subset name,psnr_masked,psnr_fg,psnr_full_image,depth_abs_fg,iou
apple,manyview_test_0,16.22478731544839,6.660985912339718,8.662890866941595,0.5735152991789598,0.0
backpack,manyview_test_0,18.664239087697137,12.092836660079621,3.9911394799946835,0.7187691122198704,0.0
ball,manyview_test_0,17.053273275949497,11.47813547143793,5.494760070704971,0.24760313752451854,0.0
banana,manyview_test_0,19.09250116156104,5.624412642679121,4.915562631182255,0.6388887597635459,0.0
baseballbat,manyview_test_0,17.662719299079523,3.56448996833759,6.856655466723437,0.5858372717711078,0.0
baseballglove,manyview_test_0,15.822024491958919,9.008496845518556,4.958078518403922,0.517665349356982,0.0
bench,manyview_test_0,16.177405149477067,5.64144135201049,6.639758049666188,0.9396015318702626,0.0
bicycle,manyview_test_0,18.929300038845177,8.384269505927424,4.978158575183426,0.7192708133061682,0.0
book,manyview_test_0,14.243260388807064,6.680398318324483,5.9082871869853735,0.9097958583065434,0.0
bottle,manyview_test_0,14.627587579689477,5.485474059329347,5.806882899714011,1.2365226740951725,0.0
bowl,manyview_test_0,12.58297015755071,4.721445807873399,6.174942733659999,0.5651215302382757,0.0
broccoli,manyview_test_0,15.348378477682894,9.138928269423888,6.406522886996562,0.46622630548488525,0.0
cake,manyview_test_0,12.406031259153915,9.13497199802905,6.954300602123617,0.7135451548332193,0.0
car,manyview_test_0,10.536444455719398,6.3033794761422826,5.589254154468083,0.6075981188742273,0.0
carrot,manyview_test_0,15.052122330808963,5.001683408210913,6.975324034802911,0.6913476205193215,0.0
cellphone,manyview_test_0,18.548592045129272,5.477199696294225,5.405821575968376,0.8925134146832333,0.0
chair,manyview_test_0,9.288750627933801,5.559044610507649,5.063084903423689,0.5832447059416495,0.0
couch,manyview_test_0,15.542901771081734,10.090205474555033,7.091879909602398,0.530379736402723,0.0
cup,manyview_test_0,14.565042555686277,4.3989084024686305,5.8416712646107225,0.9809843195171222,0.0
donut,manyview_test_0,15.455254561260311,7.186638190791148,6.08943365801032,0.42916104004956795,0.0
frisbee,manyview_test_0,16.030436839496698,8.25580372425949,3.6125508386557295,0.7820506512812717,0.0
hairdryer,manyview_test_0,22.640570140053246,11.702523731191262,4.159711019086314,0.616971255937149,0.0
handbag,manyview_test_0,24.14781075331437,15.091930028917984,5.223221264801334,0.562664145074455,0.0
hotdog,manyview_test_0,12.244917262623947,4.72460505473762,6.9914703226785,0.5147290560374835,0.0
hydrant,manyview_test_0,16.892200853920816,6.5057584631969645,6.307555495359107,0.8690763104982895,0.0
keyboard,manyview_test_0,14.937059706035933,10.816605585432766,4.857196169187754,0.5188802050007122,0.0
kite,manyview_test_0,15.068337896849323,6.205118297721433,5.276287557112783,0.7494832801627337,0.0
laptop,manyview_test_0,14.59345603707514,7.090074167371421,6.2162237610589814,0.7413216109605885,0.0
motorcycle,manyview_test_0,14.442903913583953,8.56222345535462,6.50899995433291,0.7010114811016933,0.0
mouse,manyview_test_0,29.8885518296015,14.145685466149715,5.406173914859613,0.5942925002348606,0.0
orange,manyview_test_0,11.525661011646141,5.745001890928845,5.983235030110308,0.327592487953461,0.0
parkingmeter,manyview_test_0,18.046203929985666,6.461002560728408,5.027716754597319,1.5829406195750064,0.0
pizza,manyview_test_0,15.152783189315754,6.578112135320982,7.482842326935612,0.7078538179251567,0.0
plant,manyview_test_0,20.369369422864448,11.73336728848978,5.490938199184393,0.5563616188902266,0.0
remote,manyview_test_0,21.93996425442841,9.915599775483262,3.2277628694594647,0.8952884887902877,0.0
sandwich,manyview_test_0,14.156122339232516,4.782614236412581,5.172885855269289,0.4726663784145917,0.0
skateboard,manyview_test_0,17.199716318802558,9.3986630162228,6.582697215433262,0.7526901207787688,0.0
suitcase,manyview_test_0,20.5543872349586,15.449636313939182,6.392103915747007,0.5623042520735794,0.0
teddybear,manyview_test_0,15.056483227336162,6.023824258666201,2.385989674021068,0.6859612539860361,0.0
toaster,manyview_test_0,17.538889427176077,10.389092700641873,7.350896986214959,0.6917412312874205,0.0
toilet,manyview_test_0,8.581683038527455,4.304701570881858,5.715072710684154,0.5228074506396895,0.0
toybus,manyview_test_0,13.421701717928093,5.104459961535013,7.832131890256459,0.5177220835646305,0.0
toyplane,manyview_test_0,25.939823270757692,11.015747754038403,5.005751206904976,0.5705696772343116,0.0
toytrain,manyview_test_0,17.831418296523193,7.494011795501741,4.629191510823262,0.6318052729776739,0.0
toytruck,manyview_test_0,20.369297725379987,9.285414438061778,4.844672681479939,0.48828556766453685,0.0
umbrella,manyview_test_0,12.752391495654509,6.657169727823324,2.556125460617257,0.428359657679186,0.0
vase,manyview_test_0,20.277671704818363,6.07655429478755,4.941408622390838,0.8391219139438616,0.0
wineglass,manyview_test_0,19.455250191811363,7.197566433072046,6.442702595780869,3.173690609010777,0.0
MEAN,-,16.64330518875463,7.882212795773946,5.6547484431710435,0.7209548906794958,0.0
"""
}


for task in [CO3DTask.FEW_VIEW, CO3DTask.MANY_VIEW]:
    for seq_set in [CO3DSequenceSet.DEV, CO3DSequenceSet.TEST]:
        BLANK_PREDICTION_RESULTS[(task, seq_set)] = _read_result_csv(
            CSVs[f"{task.value}_{seq_set.value}"]
        )
