notion_map = {
#    # Maintainer = Shadi
    "Armenia": "https://www.notion.so/59b0c69000274ff1b1ee83c0c6cea5a0?v=b6a8a7f3941c47ecae19ffe621482f59",
    "Austria": "https://www.notion.so/6d378966ff5c4e268376627f86f4dc2f?v=0dacaca7b49c4700a49efa550fb5f6c2",
    "Bahrain": "https://www.notion.so/9c9f2c905ae545faacaf13071d2eb033?v=7fe7cb6968bc42818030a583bac34998",
    "Colombia": "https://www.notion.so/2461e76d09cf4f05a7278a20ff1bf279?v=6233265e684a44be8a2b570b4ad7a028",
    "Costa Rica": "https://www.notion.so/1fbb092cfc1647e9a6742b54a106d13a?v=91a240e759c5406d9f2b274ad692b1ab",
    "Croatia": "https://www.notion.so/90ba2ccb229240918edce09e52ee9030?v=f28a379e24c94a3a939a7e4d891e452b",
    "Estonia": "https://www.notion.so/02011176912449ec83155233a9fbc1d7?v=7c0c40d9b4764fbd82ea999cfe529fa4",
    "Japan": "https://www.notion.so/7ac8794671464e1bbbb5b6c8ddc279f4?v=e6c39cb9ec8142939ab017c33b5aa969",
    "Jordan": "https://www.notion.so/3d290d580fa5438cbefb4f6c7d2fc375?v=fb6da0a4a57f4236b3c528fd1e9c4ad5",
    "Lebanon": "https://www.notion.so/16670cd6b914491ebfc2b488ef1b9e89?v=b0b003183b814c33830396970fac5187",
    "Philippines": "https://www.notion.so/41f843de416c4f41b9266781361e5fdb?v=842f654ed1004b6f9022421bfce6c6d6",

    # This was data that was dropped by owid in one of their commits
    "Portugal": "https://www.notion.so/95bb43ce1ae448a386db7abaafa68e77?v=2d9cf85903d34eae8b84c4072644ec92",

    # Halim
    "Argentina": "https://www.notion.so/a09b8d4dd2c44745983ab1916445ed94?v=266d332fc93543069857c1f081ea37a6",
    "Australia":  "https://www.notion.so/519b9667ab114786be8be3b1a2306c68?v=14f86820859d48be8cbc5a3199392d25",
    "Azerbaijan": "https://www.notion.so/f587f06c442b45f09c7be0ac3f58f514?v=f73a7b62fe6b4de8ac0a6f818d5a8274",
    "Bangladesh": "https://www.notion.so/35b7ee1e86b643ef8b64417112858405?v=ecbfa792b75947c6bf21501e28bf50cd",
    "Bolivia": "https://www.notion.so/9214e05323bb48f18817ed86125cf2c0?v=f3f10addeaa94d5fafbf8a5dd543ddcc",
    "Bosnia and Herzegovina": "https://www.notion.so/c4d504fb061d4a2e8d1e9b4eef16f8db?v=082817d19e1746ccbe9f98999752ddb5",
    "Chile": "https://www.notion.so/5b5ff55b44c84f09ba89d9e0bb5fd834?v=84f03e51051c41579958fc2054ff2c35",
    "Hungary": "https://www.notion.so/891f5d09165d4815834921adf73a673e?v=210fc8fb1d6a4bc297b1c5ec193ca7e1",
    "Indonesia": "https://www.notion.so/f9b8def063b14688bea0731ccb5dc046?v=464fba9d406a4d708a4c8274e44804f8",
    "New Zealand": "https://www.notion.so/66ea1fa7a90c4ff481e5c0eb6573ba67?v=a84dbac7dd2e422d93557007d1c5eaa1",
    "North Macedonia": "https://www.notion.so/7bba8f1ac34e43e19ff42f3d331c406d?v=6e85b5a16122473194ebcb7d1d3080a1",
    "Serbia": "https://www.notion.so/64ebd94f97d3440ca3431d9dd4850f76?v=05b88d514b6641d9a01da44a31372f7c",
    "Slovenia": "https://www.notion.so/200f4d913065463a92bdd8e009616af4?v=ed8aecb1b5354c6d905bfe8b5721ff90",
    "South Africa": "https://www.notion.so/e7be19dda46349188a998c7feb7c861c?v=457a2e1ffc514b4b9ebb59c626711b67",
    "Uruguay": "https://www.notion.so/d9f406529d9e4b0aa000c5950e05779f?v=67daa5dc0deb40869780016b2fa62e22",
    "Vietnam": "https://www.notion.so/c20b8220238c4fab8b183b3ee8276603?v=b4382f1d9c4e4d74b7c05b1fcd9183aa",
    #added on 2020-04-17
    "Australia – Australian Capital Territory": "https://www.notion.so/a01b1d1d3f2c461b966d7e8217bea90f?v=e1f4ed7027964f939f21952d54d76375",
    #added on 2020-04-18
    "Australia – New South Wales": "https://www.notion.so/09df7a2004c54912b2520d7a4e28e794?v=c93dc0298be346649d2cb1cb6f5e5005",
    "Australia – Northern Territory": "https://www.notion.so/c0d431e5685e418397a920ab900d6c76?v=fac2b1c5bdb648f885d23f9405916853",
    "Australia – Queensland": "https://www.notion.so/9c7e2bfac4be48e5aaf6475e01cb328d?v=f0a2353eb1264aefabe9ba496f0f7f66",
    "Australia – South Australia": "https://www.notion.so/58012bcae9604a5180d361fde6070bfc?v=e9f38fec9afe4fa6bc8a00533a6cc3c0",
    "Australia – Tasmania": "https://www.notion.so/b2b52ed9616547508d30f3add5766791?v=00b54ebdebc04d76befb370ef243f448",
    "Australia – Victoria": "https://www.notion.so/49b6a5b3a43c4717876544576e858ee7?v=f6cd19f331924b26b95f7a8521b69d59",
    "Australia – Western Australia": "https://www.notion.so/9b4c4f79de3f47b4b76ec63322def4f0?v=c1e9c91ad08f44e2860703dfdb089549",
    
    "Canada – Alberta": "https://www.notion.so/ba7e7cd3c6824e95bcfb3e1da5767192?v=561e03d80ba94978b32f519cda949a7c5",
    "Canada – British Columbia": "https://www.notion.so/ebf8e31657944d0e8d401116ccd4e927?v=433c297ed6a841b783124487bebe351f",
    "Canada – Manitoba": "https://www.notion.so/5e5c3776a2e84e4a9db44a1456ff195e?v=648ca967aead499daa37672eff40b052",
    "Canada – New Brunswick": "https://www.notion.so/49b83a8294cb4bc5aaaa8f788ca3f997?v=85f8d21012754ff19966a7526e1cfbf2",
    "Canada – Newfoundland and Labrador": "https://www.notion.so/2e848b5df8d94d31a88460a8df3b544e?v=2985bf2a5c244d668f68bee3efaf07f4",
    "Canada – Northwest Territories": "https://www.notion.so/4e00bbf18c644cf284e58b27b6f24945?v=6fe3114b96d649dba03d9b648b425ee7",
    "Canada – Nova Scotia": "https://www.notion.so/62a5b14f8f46424cad12da61b493f4ec?v=7c9f1e85340549dfadb34fbb1311af10",
    "Canada – Ontario": "https://www.notion.so/7faf169a1301432b8e9559cad1c22865?v=db11fd42916e46e298b7ece43f65ffa2",
    "Canada – Prince Edward Island": "https://www.notion.so/eadcc5c64f30495687c0eeee6b5e1377?v=259544024ed1449a9f1b074acab48270",
    "Canada – Quebec": "https://www.notion.so/e75f21ea02aa47b7b9d3195357e31caa?v=61246c07ad4f4f389f19819c88a14ae9",
    "Canada – Saskatchewan": "https://www.notion.so/ea8de2b0aa4f43c3867336a25858b663?v=221ce365021547f588860fcc2b2fbd7f",
    "Canada – Yukon": "https://www.notion.so/b3ba1b2380f14f0d837fa3a9b2c7a2c2?v=dd86dc0f36c94e33b9ed84fb6e8c3a6a",
    
    
    #added on 2020-04-21
    #students
    "Qatar": "https://www.notion.so/fbc2b37f8493465d90a0c14e9c960492?v=b1dfbbe4f9a94d4f8701b5200c75cbad",
    "Uzbekistan": "https://www.notion.so/658821c24b2643069030d316ce3329bb?v=cef868fc67bb4fcc9a7b3084a6972c32",
    "Iraq": "https://www.notion.so/a47b686a195141eea264686d5a8662d4?v=9c786876ba6c4aff902a925771f40d6a",
    "Cyprus": "https://www.notion.so/d5f09ecf0140477d895bee865b06b05c?v=2f152c89b5164944a30fb822a9a9452f",
    "San Marino": "https://www.notion.so/5b0dfd7d1462432fb8330d0f9b583e8b?v=f12c0c89e945463e8fa20cddbf85e15a",
    "Iraq": "https://www.notion.so/a47b686a195141eea264686d5a8662d4?v=9c786876ba6c4aff902a925771f40d6a",
    "France – Mayotte": "https://www.notion.so/7d5ccb5ba8b44ef4a9de0160ccb29860?v=e5ae93a6756f409594cf2b7f913e1599",
    "West Bank and Gaza": "https://www.notion.so/3db554ac61594c8ab38645830677858d?v=e891533561364ba192918278b7060ab6",
    "Mauritius":"https://www.notion.so/3a13f841e3ee4721860eeb33bd3b84c0?v=b573c84be994425d91b9dc68119771b8",
    "Morocco":"https://www.notion.so/63cd3f20bb0f4a35b605574f3235e3f7?v=7233d159d00f4304a354c9811b651330",
    "Dominican Republic":"https://www.notion.so/759c404422e744e486034e97befe5748?v=389dd435cbb849a7a9c1532c1682dcbf",
    "United Kingdom – Isle of Man":"https://www.notion.so/2b08339f35a0423b9a73853b8a08e3ca?v=af9d6eb7ebae4a658c3c9b96c14899ab",
    #HT
    "Barbados": "https://www.notion.so/73588debe9e04254bd27cc64d19a15bb?v=ff4b2c52099a4d76aa7f88223251f400",
    "Kosovo": "https://www.notion.so/6db35c5537004a36815aac221c04990c?v=f4005e62387541acb770adda09e52f09",
    "Kazakhstan":"https://www.notion.so/3a29a9eb91cb4a85a5f6a3a73c92a90c?v=bf87f2e5037c49508d098eb7194a0679",
    "Kyrgyzstan":"https://www.notion.so/8bfb2eb9576b418ab7740f54574de53a?v=de490e9e552640e6a3a012f44c1ccce0",
    
    #added in 2020-04-25
    "Saudi Arabia":"https://www.notion.so/391089760e4643678358aa9ef607be4f?v=e98f51aba3884022975f9f3e61cb8b6e",
    "Algeria":"https://www.notion.so/149253c4a0d642f0bb48ef79c094f99b?v=d273558d79954244a73f26a066fdca81",
    "Kuwait":"https://www.notion.so/a4c770fcc6794615b50fb7054ef37852?v=7adaf055ed374464a1dea61110c1011a",
    "Afghanistan":"https://www.notion.so/e8091f37822a4ed6b20efcc3ceb656a6?v=8b5c8d739118434586ad72a14da174a1",
    "Cote d'Ivoire":"https://www.notion.so/9339a6196ac8438e8e8ef9df6d2ac836?v=bf9d3184bae24d93a1978552a4ae51ae",
    "Venezuela":"https://www.notion.so/95d00b9478534f2999ff57985a77903f?v=8e2579a773024dadb85f4ee1475da9c4",
    "Georgia":"https://www.notion.so/53aef605320a44cea227bd32bc7136ab?v=cc149b2ecc3849c4bcc63ccf524afbb4",
    "United Kingdom – Channel Islands":"https://www.notion.so/8ad5e1613099456fab7ff7d8616bb4b5?v=34b492262de1482d8346fe260814b77c",
    
    #added in 2020-04-27
    "Denmark – Faroe Islands":"https://www.notion.so/4572876b326e4671b0ca8be614e3e5d7?v=9476926efadd4f1d9cf193c8fae969b0",
    
    #added in 2020-04-28
    "Guyana":"https://www.notion.so/00b15b1cfaac40bbab115f8bc22a02b3?v=1bf323b9374044d98ae6a6e2f6252f37",
    "United Kingdom – Bermuda":"https://www.notion.so/eae28caf7bb04e06ab7c425567fffd58?v=be204ee5e82a4ab98ef398daff47bd96",
    "Netherlands – Aruba ":"https://www.notion.so/8108ba3f33b1481381b450e4d276a308?v=3896389fad144150a099232c30b6485e",
    "Cambodia":"https://www.notion.so/9299191b17a746d4b81fd27dcc11e7a1?v=a9253a72a8ff418cbbf53e13bf4f6ce4",
    "United Kingdom – Gibraltar":"https://www.notion.so/a0077de1de9a4bf9ab99e75479921cf9?v=a82972cb04774b76bee8e9921c842ddb",
    "France – Martinique":"https://www.notion.so/705401eec6be49c8a146aa1b2e3d7626?v=cd9e38d58f6f48938ae1ae8145e88ada",
    
    #added in 2020-04-30
    "Bulgaria":"https://www.notion.so/09504393f2df4b93bf5d9a8a4414791e?v=7baa0d43b7de4a60bb27f5939a3b1b0e",
    "Malta":"https://www.notion.so/077758ffdd184c50a45e276d67c4c5f1?v=abf9ee814dcd4f2b8d3aa77cc1cc0da6",
    "Montenegro":"https://www.notion.so/f97ef019adab47a6ab06761ebd80c153?v=93b48999c450468688d058a7cc1964f9",
    "Rwanda":"https://www.notion.so/14edfcd8ff15418c9c39a66b33833ff4?v=5043dc278a4f492189fa77bec58a64db",
    "Togo":"https://www.notion.so/0e1839209d964cd59b77912cc121414e?v=6c0f00d5c139413f8f41bb9cf956fb12",
    "Burma":"https://www.notion.so/54048bda9f1a43b68ff5730cb46068ab?v=970181c8096a4e9998a0e7c48b94f307",
    
    #added in 2020-05-01
    "France – Guadeloupe":"https://www.notion.so/c00bf4bac08e448ba9f6c18aae2602e3?v=19f396ece82c4920b743cbf687ed7692",
    "United Kingdom – Cayman Islands":"https://www.notion.so/38d75cfb2e244864bbad129d9c1bb351?v=6afa466f971445f2b5abc6bf5612c74a",
    "Zambia":"https://www.notion.so/895d0f2ba719431ca74abf525f9e6eae?v=98e06d80eb3049a4a98f2b146e2ad15c",
    "France – French Polynesia":"https://www.notion.so/324a0640ffdb45079f1e10a7538d662c?v=f13fa4692f054eafa986682b333be8ee",
    "Netherlands – Sint Maarten":"https://www.notion.so/432baba684a64e80aac8acb28c03d1b9?v=6a11dfc353344c58b45bd6659568b727",
    "Brunei":"https://www.notion.so/8e51e427687045708e29e2cb1df59e3d?v=20049dd68d264743ad8a7aadd258ac9f",

    # 2020-05-11
    "Brazil": "https://www.notion.so/086536d4083d41d1af1ddb4af92ed08f?v=9b9987e7c1ab4bd39fb6e48877224e8f",
    
    #2020-05-15
    "Liechtenstein":"https://www.notion.so/dfe83dcaee8f43e7a7c7a46257bfe0f3?v=aa2d86952bf54d3da24c9368a0d613e7",
    
}


from notion.collection import NotionDate
import pandas as pd


def get_table(client, country_name):
    if not country_name in notion_map.keys():
      raise Exception(f"Country %{country_name} not found in notion map")

    page = client.get_block(notion_map[country_name])
    coll=page.collection.get()

    # column names, eg ['deaths', 'treatment', 'confirmed', 'negative', 'recoveries', 'Date', 'ID']
    colnames = [x['name'] for i,x in coll['schema'].items()]
    
    # prep for dataframe
    rows = page.collection.get_rows()
    df_all = []
    for i,r in enumerate(rows):
      df_new = {}
      for c in colnames:
        try:
          v = getattr(r,c)
        except Exception as e:
          import click
          click.secho(f"Error in country {country_name}, row {i}, column {c}", fg="red")
          click.secho(str(e), fg="red")
          import sys
          sys.exit(1)

        if type(v)==NotionDate: v = v.start
        df_new[c] = v

      df_all.append(df_new)

    # actual dataframe
    df_all = pd.DataFrame(df_all)

    return df_all


def postprocess_table(country_name, df_single):
  # Update 2020-04-09: in general, do not replace NA with 0 with fillna(0)

  if country_name=="Argentina":
    # df_single["total_cumul"] = df_single[["neg cumul labs", "neg cumul epid", "confirmed cumulative"]].fillna(0).apply(sum, axis=1)
    # Update 2020-04-09: do not use neg cumul epid
    df_single.loc[df_single["total_cumul"].isnull(),'total_cumul'] = df_single[["neg cumul labs", "confirmed cumulative"]].apply(sum, axis=1)
    return df_single
  
  #if country_name=="Armenia":
    # update 2020-04-20: turns out that the column name called "negative" is actually "total_cumul", so no need for postprocessing anymore
    #df_single["total_cumul"] = df_single[["confirmed", "negative"]].apply(sum, axis=1)
    #return df_single

  if country_name=="Cambodia":
    df_single["total_cumul"] = df_single["total_cumul_persons"]
    return df_single

  if country_name=="Canada – Saskatchewan":
    df_single["total_cumul"] = df_single[["confirmed", "negative"]].apply(sum, axis=1)
    return df_single
    
  if country_name=="United Kingdom – Channel Islands":
    df_single["total_cumul"] = df_single[["total_cumul_guernsey", "total_negative_jersey","total_positive_jersey"]].apply(sum, axis=1)
    return df_single

  if country_name=="Canada – Prince Edward Island":
    df_single["total_cumul"] = df_single[["confirmed", "negative"]].apply(sum, axis=1)
    return df_single

  if country_name=="Cyprus":
    df_single["total_cumul"] = df_single[["total_cumul_north", "total_cumul_south"]].apply(sum, axis=1)
    return df_single 
 
  if country_name=="Dominican Republic":
    df_single["total_cumul"] = df_single["PCR_tests"]
    return df_single

  if country_name=="Morocco":
    df_single["total_cumul"] = df_single[["confirmed", "negative"]].apply(sum, axis=1)
    return df_single

  if country_name=="Bolivia":
    df_single["total_cumul"] = df_single[["cumulative negative", "cumulative confirmed"]].apply(sum, axis=1)
    return df_single

  # Update 2020-04-10: turns out they have a total cumulative tests field
  # Update 2020-04-14: Halim dropped the total field and is just recording positive and negative
  if country_name=="Bosnia and Herzegovina":
    df_single["total_cumul"] = df_single[["negative_cumul", "positive_cumul"]].apply(sum, axis=1)
    return df_single

  #Updated 2020-05-01 Halim added total_cumul from source 
  #if country_name=="Colombia":
  #  df_single["total_cumul"] = df_single[["confirmed cumul", "negative cumul"]].apply(sum, axis=1)
  #  return df_single

  if country_name=="Costa Rica":
    df_single["total_cumul"] = df_single[["confirmed cumul", "negative cumul"]].apply(sum, axis=1)
    return df_single

  #if country_name=="Croatia":
  #  # Croatia has a total_cumul column, but for some dates that are missing it, we can calculate it from confirmed + negative
  #  df_single["total_cumul"] = df_single[["confirmed cumul", "negative cumul"]].fillna(0).apply(sum, axis=1)
  #  return df_single

  # Update 2020-04-09: just use the total_cumul field directly
  #if country_name=="Estonia":
  #  df_single["total_cumul"] = df_single[["confirmed", "negative"]].apply(sum, axis=1)
  #  return df_single

  if country_name=="Lebanon":
    df_single["total_cumul"] = df_single[["total_cumul_labs", "total_cumul_airport"]].apply(sum, axis=1)
    return df_single

  if country_name=="Philippines":
    df_single["total_cumul"] = df_single[["confirmed_cumul", "negative_cumul","total individuals tested"]].apply(sum, axis=1)
    return df_single

  if country_name=="San Marino":
    df_single["total_cumul"] = df_single[pd.notnull(df_single.total_cumul_swabs)][["total_cumul_serological", "total_cumul_swabs"]].fillna(0).apply(sum, axis=1)
    return df_single

  if country_name=="Vietnam":
    df_single["total_cumul"] = df_single[["total positive", "total negative"]].apply(sum, axis=1)
    return df_single

  #print(f"no postprocessing for country {country_name}")
  return df_single



import pandas as pd
from notion.client import NotionClient

class L0ImportBiominers:

  def get_notion_token(self, notion_fn):
    token_v2=open(notion_fn,"r").read().strip()
    self.client = NotionClient(token_v2=token_v2)

  def fetch_tables(self):
    # get all
    print("fetching all tables")
    df_global = []
    for country_name in sorted(list(notion_map.keys())):
      #if country_name<"Bulgaria": continue # FIXME

      print(f"Getting table for {country_name}")
      df_single = get_table(self.client, country_name)
      print("%s .. %s"%(country_name, df_single.shape))

      df_single = postprocess_table(country_name, df_single)
      # df_single = df_single[["Date","total_cumul"]]
      #if country_name=="Lebanon":
      #  import pdb
      #  pdb.set_trace()
      # df_single.head(1).transpose()
        
      # At this point, all tables should have a total_cumul field and a Date field
      # whether it comes from the notion.so table directly, or from the postprocess_table function
      assert "Date" in df_single.columns
      assert "total_cumul" in df_single.columns

      df_single["country_t11"] = country_name
      df_global.append(df_single)

    # merge
    df_global = pd.concat(df_global, axis=0)
    df_global = df_global.loc[pd.notnull(df_global.total_cumul),]
    df_global["total_cumul"] = df_global.total_cumul.astype(int)
    df_global = df_global.sort_values(["country_t11", "Date"], ascending=True)

    if df_global[["country_t11","Date"]].duplicated().any():
      df_dupes = df_global[df_global[["country_t11","Date"]].duplicated()]
      print(df_dupes[~df_dupes.duplicated()][["country_t11","Date"]])
      import click
      click.secho("Found %i duplicates in biominers data"%(df_global[["country_t11","Date"]].duplicated().sum()), fg="red")
      import sys
      sys.exit(3)

    self.df_global = df_global

  def to_csv(self, fn_biominers):
    # save
    # fn_biominers = "multiple-biominers-gitrepo.csv"
    df_save = self.df_global[["country_t11", "Date", "total_cumul"]]
    df_save.to_csv(fn_biominers, index=False)


