import os, gzip, json
import numpy as np
from pandas import DataFrame
from mpcontribs.io.core.recdict import RecursiveDict
from mpcontribs.io.core.utils import get_composition_from_string
from mpcontribs.users.boltztrap.rest.rester import BoltztrapRester

def run(mpfile, nmax=5, dup_check_test_site=True):

    # book-keeping
    existing_mpids = {}
    for b in [False, True]:
        with BoltztrapRester(test_site=b) as mpr:
            for doc in mpr.query_contributions(criteria=mpr.query):
                existing_mpids[doc['mp_cat_id']] = doc['_id']
        if not dup_check_test_site:
            break

    # extract data from json files
    keys = ['pretty_formula', 'volume']
    input_dir = mpfile.hdata.general['input_dir']
    for idx, fn in enumerate(os.listdir(input_dir)):
        print(fn)
        input_file = gzip.open(os.path.join(input_dir, fn), 'rb')
        try:
            data = json.loads(input_file.read())

            # add hierarchical data (nested key-values)
            # TODO: extreme values for Seebeck, conductivity, power factor, zT, effective mass
            mpfile.add_hierarchical_data(
                RecursiveDict((k, data[k]) for k in keys),
                identifier=data['mp_id']
            )

            # add structure; TODO necessary only if not canonical from MP
            #name = get_composition_from_string(data['pretty_formula'])
            #mpfile.add_structure(
            #    data['cif_structure'], name=name,
            #    identifier=data['mp_id'], fmt='cif'
            #)

            # add data table
            for doping_type in ['n', 'p']:
                seebeck = data['GGA']['seebeck_doping'][doping_type]
                seebeck_averages, dopings, columns = [], None, ['T']
                temps = sorted(map(int, seebeck.keys()))
                for temp in temps:
                    row = ['{} K'.format(temp)]
                    if dopings is None:
                        dopings = sorted(map(float, seebeck[str(temp)].keys()))
                    for doping in dopings:
                        doping_str = '%.0e' % doping
                        if len(columns) <= len(dopings):
                            columns.append(doping_str + ' cm-3')
                        eigs = seebeck[str(temp)][doping_str]['eigs']
                        row.append(np.mean(eigs))
                    seebeck_averages.append(row)
                df = DataFrame.from_records(seebeck_averages, columns=columns)
                table_name = doping_type + '-type average seeback'
                mpfile.add_data_table(data['mp_id'], df, table_name)

        finally:
            input_file.close()
        if idx >= nmax+1:
            break

@@ -274,9 +274,9 @@
 max=
-5
+1
 , du
@@ -1570,16 +1570,83 @@
 a table%0A
+            for prop_name in %5B'seebeck_doping','cond_doping'%5D:%0A    
         
@@ -1696,23 +1696,24 @@
         
-seebeck
+    prop
  = data%5B
@@ -1719,32 +1719,25 @@
 %5B'GGA'%5D%5B
-'seebeck_doping'
+prop_name
 %5D%5Bdoping
@@ -1751,39 +1751,40 @@
                 
-seebeck
+    prop
 _averages, dopin
@@ -1829,16 +1829,20 @@
         
+    
 temps = 
@@ -1857,23 +1857,20 @@
 ap(int, 
-seebeck
+prop
 .keys())
@@ -1863,32 +1863,36 @@
 , prop.keys()))%0A
+    
                 
@@ -1930,16 +1930,20 @@
         
+    
 row = %5B'
@@ -1982,16 +1982,20 @@
         
+  
+  
 if dopin
@@ -2006,16 +2006,20 @@
 s None:%0A
+    
         
@@ -2062,23 +2062,20 @@
 (float, 
-seebeck
+prop
 %5Bstr(tem
@@ -2079,32 +2079,36 @@
 temp)%5D.keys()))%0A
+    
                 
@@ -2154,24 +2154,28 @@
             
+   
+ 
 doping_str =
@@ -2187,24 +2187,28 @@
 e' %25 doping%0A
+    
             
@@ -2280,16 +2280,20 @@
         
+   
+ 
 columns.
@@ -2349,22 +2349,23 @@
     
+    
 eigs = 
-seebeck
+prop
 %5Bstr
@@ -2392,16 +2392,20 @@
 'eigs'%5D%0A
+    
         
@@ -2466,23 +2466,24 @@
         
-seebeck
+    prop
 _average
@@ -2496,16 +2496,20 @@
 nd(row)%0A
+    
         
@@ -2544,23 +2544,20 @@
 records(
-seebeck
+prop
 _average
@@ -2576,16 +2576,20 @@
 olumns)%0A
+    
         
@@ -2642,17 +2642,26 @@
 age 
-seeback'%0A
+' + prop_name%0A    
     

add new property