import os
import rasterio

import numpy as np


class wp_demographics(object):
    def __init__(self, wp_file):
        self.wp_file = wp_file
        self.age = os.path.basename(wp_file).split("_")[2]
        self.gender = os.path.basename(wp_file).split("_")[1]
        self.country = os.path.basename(wp_file).split("_")[0]
        self.raster = rasterio.open(wp_file)

    def __str__(self):
        return f"{self.country} - {self.gender} - {self.age}"


class wp_vulnerability(object):
    def __init__(self, wp_files, vulnerability_definition):
        """Process the wp demographic information into a single map of vulnerability

        INPUT
        wp_files [list of wp_demographics]
        vulnerability [dictionary of definitions]
        """
        self.wp_files = wp_files
        self.v_def = vulnerability_definition
        # classify wp_files into categories
        vul_files = {}
        for key, value in vulnerability_definition.items():
            min_val = int(key.split("-")[0])
            max_val = int(key.split("-")[1])
            for f in wp_files:
                # print(f'{f.age}: {min_val}')
                if (int(f.age) > min_val) and (int(f.age) <= max_val):
                    try:
                        vul_files[key].append(f)
                    except:
                        vul_files[key] = [f]
        self.vul_files = vul_files

    def calculate_total_pop(self):
        vul_results = {}
        for key, values in self.vul_files.items():
            for wp in values:
                curR = wp.raster.read()
                try:
                    final = final + curR
                except:
                    final = curR
        self.total_pop = final

    def calculate_vulnerability(self):
        vul_results = {}
        for key, values in self.vul_files.items():
            vul_value = self.v_def[key]
            for wp in values:
                curR = wp.raster.read()
                curR = curR * vul_value
                try:
                    vul_results[key] = vul_results[key] + curR
                except:
                    vul_results[key] = curR
        self.vul_res = vul_results

    def combine_results(self, out_vulnerability="", out_pop=""):
        for key, values in self.vul_res.items():
            for v in values:
                try:
                    final = final + v
                except:
                    final = v
        final[final < 0] = np.nan
        if out_vulnerability != "":
            with rasterio.open(
                out_vulnerability, "w", **self.wp_files[0].raster.meta
            ) as outR:
                outR.write_band(1, final)
        else:
            return final

        if out_pop != "":
            with rasterio.open(out_pop, "w", **self.wp_files[0].raster.meta) as outR:
                outR.write(self.total_pop)
