#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Lectura directa de archivo MRAMS (.ctl + .dat) sin PyGrADS.
Extrae TEMP (tempk) en 3 niveles (ej. 14.5, 40, 100 m).
Compatible con Python 2.7
"""

import struct

# =============================
# CONFIGURACIÓN DEL USUARIO
# =============================

ctl_file = "Jezero_Mars2020_Ls0-atm-S-g07.ctl"
dat_file = ctl_file.replace(".ctl", ".dat")

# Índices del punto de Jezero (según tu Fortran)
ix = 43 - 1   # Fortran empieza en 1, Python en 0
iy = 93 - 1

# Niveles deseados (Fortran -> Python índices)
levels = [2, 3, 4]  # 2≈14.5m, 3≈40m, 4≈100m
levels = [l-1 for l in levels]

# Rango de tiempos
t_start = 290 - 1
t_end   = 579 - 1

# Variable a extraer
varname = "tempk"

# Archivo de salida
output_txt = "MRAMS_Jezero_tempk_levels.txt"

# =============================
# PARSEAR EL CTL
# =============================

nx, ny, nz, nt = None, None, None, None
vars_info = []
inside_vars = False

with open(ctl_file) as f:
    for line in f:
        parts = line.strip().split()
        if len(parts) == 0:
            continue
        key = parts[0].lower()
        if key == "xdef":
            nx = int(parts[1])
        elif key == "ydef":
            ny = int(parts[1])
        elif key == "zdef":
            nz = int(parts[1])
        elif key == "tdef":
            nt = int(parts[1])
        elif key == "vars":
            inside_vars = True
            continue
        elif key == "endvars":
            inside_vars = False
            continue
        elif inside_vars:
            vname = parts[0]
            vlevs = int(parts[1])
            vars_info.append((vname, vlevs))

# Buscar variable
offset_var = 0
target_var = None
for vname, vlevs in vars_info:
    if vname == varname:
        target_var = (offset_var, vlevs)
        break
    offset_var += vlevs

if target_var is None:
    raise Exception("Variable {} no encontrada en el ctl".format(varname))

var_offset, var_levels = target_var

print("Archivo con dimensiones: nx={}, ny={}, nz={}, nt={}".format(nx, ny, nz, nt))
print("Variable {} encontrada con {} niveles".format(varname, var_levels))

# =============================
# LECTURA DIRECTA DEL .DAT
# =============================

field_size = nx * ny  # puntos horizontales
fdat = open(dat_file, "rb")

with open(output_txt, "w") as fout:
    fout.write("time_index " + " ".join(["{}_level{}".format(varname, l+1) for l in levels]) + "\n")

    for tt in range(t_start, t_end+1):
        row = [str(tt+1)]
        for lev in levels:
            pos = (tt * sum(v[1] for v in vars_info) * field_size +   # tiempos previos
                   var_offset * field_size +                          # variables previas
                   lev * field_size +                                 # niveles previos
                   (iy * nx + ix))                                    # punto (x,y)

            fdat.seek(pos * 4)  # float32 (4 bytes)
            val = struct.unpack("f", fdat.read(4))[0]
            row.append(str(val))
        fout.write(" ".join(row) + "\n")

fdat.close()
print("Archivo guardado:", output_txt)