import gzip
import struct
import numpy as np

def reader(file_path):
    with gzip.open(file_path, 'rb') as f:
        # unpack header
        fmt1 = '<9i4s10i'
        size1 = struct.calcsize(fmt1)
        data1 = struct.unpack(fmt1, f.read(size1))

        yyyy, mm, dd, hh, mn, ss, nx, ny, nz = data1[0:9]
        proj = data1[9].decode('ascii').strip()
        map_scale, projlat1, projlat2, projlon, alon, alat, xy_scale, dx, dy, dxy_scale = data1[10:20]

        fmt2 = f'<{nz}i'
        zht = list(struct.unpack(fmt2, f.read(struct.calcsize(fmt2))))
        
        fmt3 = '<11i'
        size3 = struct.calcsize(fmt3)
        data3 = struct.unpack(fmt3, f.read(size3))
        z_scale, i_bb_mode = data3[0:2]
        unkn01 = data3[2:11]

        varname = f.read(20).decode('ascii').strip()
        varunit = f.read(6).decode('ascii').strip()
        fmt4 = '<3i'
        var_scale, missing, nradar = struct.unpack(fmt4, f.read(struct.calcsize(fmt4)))

#        print(yyyy, mm, dd, hh, mn, ss, nx, ny, nz)
#        print(proj)
#        print(map_scale, projlat1, projlat2, projlon, alon, alat, xy_scale, dx, dy, dxy_scale)
#        print(zht)
#        print(z_scale, i_bb_mode)
#        print(unkn01)
#        print(varname)
#        print(varunit)
#        print(var_scale, missing, nradar)
        mosradar = []
        for i in range(nradar):
            raderid = f.read(4).decode('ascii', errors='ignore').strip()
#            print(raderid)
            mosradar.append(raderid)

        # reading values
        f_ungzip = np.frombuffer(f.read(), dtype=np.int16)
        phys = f_ungzip
#        print(phys.shape)
#        print(nx, ny, nz)
        phys = np.reshape(phys, (ny, nx, nz)) / var_scale
#        import matplotlib.pyplot as plt
#        phys[phys==-999.0]=np.nan
#        plt.imshow(phys[:, :, 0], vmin=0, vmax=60, cmap='jet')
#        plt.savefig('tmp.png', dpi=200)

    return phys, missing

def get_lonlat_z(file_path):
    with gzip.open(file_path, 'rb') as f:
                # unpack header
        fmt1 = '<9i4s10i'
        size1 = struct.calcsize(fmt1)
        data1 = struct.unpack(fmt1, f.read(size1))

        yyyy, mm, dd, hh, mn, ss, nx, ny, nz = data1[0:9]
        proj = data1[9].decode('ascii').strip()
        map_scale, projlat1, projlat2, projlon, alon, alat, xy_scale, dx, dy, dxy_scale = data1[10:20]

        fmt2 = f'<{nz}i'
        zht = list(struct.unpack(fmt2, f.read(struct.calcsize(fmt2))))

        fmt3 = f'<1i'
        zscale = struct.unpack(fmt3, f.read(struct.calcsize(fmt3)))

        lonmin = alon / xy_scale
        lonmax = lonmin + dx / dxy_scale * (nx - 1)
        latmax = alat / xy_scale
        latmin = latmax - dy / dxy_scale * (ny - 1)

        lon = np.linspace(lonmin, lonmax, num=nx, endpoint=True)
        lat = np.linspace(latmin, latmax, num=ny, endpoint=True)
        
        mlon, mlat = np.meshgrid(lon, lat)

        z = np.array(zht) / zscale

    return mlon, mlat, z

def test_main():
    phys, missing = reader("20170601/cref_mosaic/data/ASIA_CREF.20170601.000000.gz")
    #phys, missing = reader("/data3/jinni/mrefl_mosaic/2023/202301/20230101/MREF3D21L.20230101.2300.gz")
    #phys, missing = reader("/data3/jinni/compref_mosaic/2023/202301/20230101/COMPREF.20230101.0000.gz")
    mlon, mlat, z = get_lonlat_z("/data3/jinni/mrefl_mosaic/2023/202301/20230101/MREF3D21L.20230101.2300.gz")

if __name__ == "__main__":
    test_main()