#!/usr/bin/env python
# -*- coding: utf-8 -*-

# ==============================================================================
# author          :Ghislain Vieilledent
# email           :ghislain.vieilledent@cirad.fr, ghislainv@mtloz.fr
# web             :https://ecology.ghislainv.fr
# python_version  :>=3.6
# license         :GPLv3
# ==============================================================================

import os
from typing import Dict, List

import terracotta as tc

# Define the location of the SQLite database
# (this will be created if it doesn't already exist)
DB_NAME = "./terracotta.sqlite"

# Define the list of keys that will be used to identify datasets.
# (these must match the keys of the "key_values" dicts defined in
# RASTER_FILES)
KEYS = ["name", "year", "cont"]

# Define a list of raster files to import
# (this is a list of dictionaries, each with a file path and the
# values for each key - make sure the order matches the order of
# KEYS defined above)
#
# This part of the script could be replaced with something that
# makes sense for your data - it could use a glob expression to
# find all TIFFs and a regular expression pattern to extract the
# key values, or it could read from a CSV, or use some other
# reference or metadata generating process.
RASTER_FILES_1 = [
    {
        "key_values": {
            "name": name,
            "year": year,
            "cont": cont,
        },
        "path": f"./{name}_{year}_{cont}_merc.tif",
    } for name in ["fcc"] for year in ["123", "2050", "2100"] for cont in ["AME", "AFR", "ASI"]
]
RASTER_FILES_2 = [
    {
        "key_values": {
            "name": name,
            "year": year,
            "cont": cont,
        },
        "path": f"./{name}_{year}_{cont}_merc.tif",
    } for name in ["prob"] for year in ["2020"] for cont in ["AME", "AFR", "ASI"]
]
RASTER_FILES = RASTER_FILES_1 + RASTER_FILES_2

def load(db_name: str, keys: List[str], raster_files: List[Dict]):
    # get a TerracottaDriver that we can use to interact with
    # the database
    driver = tc.get_driver(db_name)

    # create the database file if it doesn't exist already
    if not os.path.isfile(db_name):
        driver.create(keys)

    # check that the database has the same keys that we want
    # to load
    assert list(driver.key_names) == keys, (driver.key_names, keys)

    # connect to the database
    with driver.connect():
        # insert metadata for each raster into the database
        for raster in raster_files:
            md = driver.compute_metadata(raster["path"],
                                         use_chunks=None,
                                         max_shape=(1024, 1024))
            driver.insert(raster["key_values"], raster["path"],
                          metadata=None, skip_metadata=True)
            
if __name__ == "__main__":
    load(DB_NAME, KEYS, RASTER_FILES)


# Tests
FILE = [RASTER_FILES[len(RASTER_FILES) - 1]]
load(DB_NAME, KEYS, FILE)

driver = tc.get_driver(DB_NAME)
print(driver.get_metadata(["prob", "2020", "ASI"]))

    
# End of file