import netCDF4 as nc
from pathlib import Path
import os
import numpy as np
import logging

# setup logging
logger = logging.getLogger()
logging.basicConfig()
logger.setLevel("DEBUG")

file = '/data/daves/AR_grids/test/processed_winds_d02_2019-02-08_12_00_00.nc'
datetime = os.path.basename(file)[20:].split('.')[0] # 2019-02-08_18:00:00

dir_out = f'/home/sreiner/AR_grids/binary_files_test/{datetime[:10]}'
logging.debug(f'input file: {file}')
logging.debug(f'output directory: {dir_out}')

Path(dir_out).mkdir(parents=True, exist_ok=True) # create nested directory for output files

def calc_splice(x,y):
    modx = x % 8
    if modx == 0 or modx == 7:
        retx = 0
    else:
        retx = modx
    mody = y % 8
    if mody == 0 or mody == 7:
        rety = 0
    else:
        rety = mody
    return retx, rety

with nc.Dataset(file, 'r') as f:
    logging.info('writing xlat')
    with open(f'{dir_out}/{datetime}_xlat.bin', 'wb') as wb:
        data = f['XLAT'][0]
        x, y = np.shape(data)
        logging.debug(f'size before: ({x}, {y})')
        xp, yp = calc_splice(x, y)
        d = np.ascontiguousarray(data[xp::2, yp::2])
        logging.debug(f'size after: {np.shape(d)}')
        wb.write(d)

    logging.info('writing xlong')
    with open(f'{dir_out}/{datetime}_xlong.bin', 'wb') as wb:
        data = f['XLONG'][0]
        if np.shape(data) != (x, y):
            x, y = np.shape(data)
            xp, yp = calc_splice(x, y)
        d = np.ascontiguousarray(data[xp::2, yp::2])
        wb.write(d)

    # umet, vmet
    levels = [55, 59, 63, 70, 75, 80, 87, 91]
    pressure = [200, 250, 300, 400, 500, 600, 700, 850]
    for lvl, p in zip(levels, pressure):
        umet = f['umet'][0, lvl]  # read slice of umet
        vmet = f['vmet'][0, lvl]  # read slice of vmet
        logging.info(f'writing pressure {p}')

        # resample then calculate mins/maxs or calculate min/max then resample? Try first one
        if np.shape(umet) != (x,y):
            x, y = np.shape(umet)
            xp, yp = calc_splice(x, y)
        umet = np.ascontiguousarray(umet[xp::2, yp::2])

        if np.shape(vmet) != (x,y):
            x, y = np.shape(vmet)
            xp, yp = calc_splice(x, y)
        vmet = np.ascontiguousarray(vmet[xp::2, yp::2])

        logging.debug(f'shape umet: {np.shape(umet)}')
        logging.debug(f'shape vmet: {np.shape(vmet)}')
        with open(f'{dir_out}/{datetime}_prs_umet_{p}.bin', 'wb') as wb:
            wb.write(umet)

        with open(f'{dir_out}/{datetime}_prs_vmet_{p}.bin', 'wb') as wb:
            wb.write(vmet)