""" handle GMF's modifier files that update model inputs with meta- parameters

- reads rocktype JSON and update model input
- reads upflow JSON and update model input
"""

import json
import os
import sys

def compute_rocktypes(rocktype_json, warn=True):
    """ compute the derived rocktype properties and update in place

    if warn, will raise exception when rules are not followed:
        - along modifier >= 1.0
        - across modifier <= 1.0

    Example entry in rocktype_json:
    - base rocktype
        "A0000": {
          "permeability": [1e-11, 1e-11, 1e-12 ],
          "porosity": [0.25 ]
        },

    - fault modified rocktype
        "AA000": {
          "base": ["A0000"],
          "fault direction": [1],
          "modifier": {
            "along": [1],
            "across": [1],
            "vertical": [1],
            "porosity": [1]
          },
          "permeability": [1e-11, 1e-11, 1e-12],
          "porosity": [0.25]
        },

    - fault interesection rocktype
        "SVU11": {
          "base": ["S0011"],
          "modifier": {
            "k1": [2.5],
            "k2": [0.4],
            "k3": [5],
            "porosity": [1]
          },
          "permeability": [2.5000000000000004e-15, 4.0000000000000004e-16, 5.000000000000001e-15],
          "porosity": [0.1]
        }

    """
    for rt_name, rt in rocktype_json.items():
        if 'base' not in rt:
            # base rocktype
            pass
        else:
            if 'fault direction' in rt:
                fault_dir = rt['fault direction'][0]
                if float(rt['modifier']['along'][0]) < 1.0:
                    msg = f"Rocktype {rt_name} has fault modifier 'along' < 1.0"
                    if warn: raise Exception(msg)
                    else: print("Warning: " + msg)
                if float(rt['modifier']['across'][0]) > 1.0:
                    msg = f"Rocktype {rt_name} has fault modifier 'across' > 1.0"
                    if warn: raise Exception(msg)
                    else: print("Warning: " + msg)
                if fault_dir == 1:
                    k_modifier = [
                        float(rt['modifier']['along'][0]),
                        float(rt['modifier']['across'][0]),
                        float(rt['modifier']['vertical'][0]),
                    ]
                elif fault_dir == 2:
                    k_modifier = [
                        float(rt['modifier']['across'][0]),
                        float(rt['modifier']['along'][0]),
                        float(rt['modifier']['vertical'][0]),
                    ]
                else:
                    raise Exception(f"Fault direction {rt['fault direction'][0]} not 1 or 2 in rocktype {rt_name}")
            else:
                k_modifier = [
                    float(rt['modifier']['k1'][0]),
                    float(rt['modifier']['k2'][0]),
                    float(rt['modifier']['k3'][0]),
                ]
            c_modifier = 1.0
            if 'conductivity' in rt['modifier']:
                c_modifier = float(rt['modifier']['conductivity'][0])
            # update values in place
            base_name = rt['base'][0]
            if base_name in rocktype_json:
                rt['permeability'] = [k * m for k, m in zip(
                    rocktype_json[base_name]['permeability'], k_modifier)]
                rt['porosity'][0] = rocktype_json[base_name]['porosity'][0] * rt['modifier']['porosity'][0]
                if 'conductivity' in rocktype_json[base_name]:
                    # print(f"Rocktype {rt_name} use conductivity from base rocktype {base_name}")
                    rt['conductivity'] = [rocktype_json[base_name]['conductivity'][0] * c_modifier]
                else:
                    if 'conductivity' in rt:
                        print(f"(gimu)    Rocktype {rt_name} remove conductivity because no conductivity in base rocktype {base_name}")
                        del rt['conductivity']
            else:
                msg = f"Base rocktype {base_name} not found for rocktype {rt_name}"
                if warn:
                    raise Exception(msg)
                else:
                    print("(gimu)    Warning: " + msg)
                    rt['error'] = msg


def update_waiwera_model_rocktypes(wjson, rtjson, skip_missing=True):
    """ rtjson is loaded from the GMF modifier _rocktypes.json
    """
    for rt in wjson['rock']['types']:
        if rt['name'] not in rtjson:
            if skip_missing:
                print(f"(gimu)    Rocktype {rt['name']} not found in _rocktypes.json... skipped")
                continue
            else:
                raise Exception(f"(gimu)    Rocktype {rt['name']} not found in _rocktypes.json")
        modifier = rtjson[rt['name']]
        rt['permeability'] = modifier['permeability']
        rt['porosity'] = modifier['porosity'][0]
        if 'conductivity' in modifier:
            rt['wet_conductivity'] = modifier['conductivity'][0]
            rt['dry_conductivity'] = modifier['conductivity'][0]

def update_aut2_model_rocktypes(dat, rtjson, skip_missing=True):
    """ rtjson is loaded from the GMF modifier _rocktypes.json
    """
    for rt in dat.grid.rocktypelist:
        if rt.name not in rtjson:
            if skip_missing:
                print(f"(gimu)    Rocktype {rt.name} not found in _rocktypes.json... skipped")
                continue
            else:
                raise Exception(f"(gimu)    Rocktype {rt.name} not found in _rocktypes.json")
        modifier = rtjson[rt.name]
        rt.permeability = modifier['permeability']
        rt.porosity = modifier['porosity'][0]
        if 'conductivity' in modifier:
            rt.conductivity = modifier['conductivity'][0]

