import time
import math
from machine import I2C, Pin, ADC

# ── I2C Setup ──────────────────────────────
i2c = I2C(0, sda=Pin(4), scl=Pin(5), freq=400000)
adc = ADC(26)

# ── SSD1306 OLED Driver (no library needed) ─
OLED_ADDR = 0x3C

def oled_cmd(cmd):
    i2c.writeto(OLED_ADDR, bytes([0x00, cmd]))

def oled_init():
    for cmd in [
        0xAE, 0x20, 0x00, 0xB0, 0xC8,
        0x00, 0x10, 0x40, 0x81, 0xFF,
        0xA1, 0xA6, 0xA8, 0x3F, 0xA4,
        0xD3, 0x00, 0xD5, 0xF0, 0xD9,
        0x22, 0xDA, 0x12, 0xDB, 0x20,
        0x8D, 0x14, 0xAF
    ]:
        oled_cmd(cmd)

def oled_clear():
    for page in range(8):
        oled_cmd(0xB0 + page)
        oled_cmd(0x00)
        oled_cmd(0x10)
        i2c.writeto(OLED_ADDR, bytes([0x40] + [0x00] * 128))

# Font 5x7 basic characters
FONT = {
    ' ': [0,0,0,0,0],
    'A': [0x7E,0x11,0x11,0x11,0x7E],
    'B': [0x7F,0x49,0x49,0x49,0x36],
    'C': [0x3E,0x41,0x41,0x41,0x22],
    'D': [0x7F,0x41,0x41,0x22,0x1C],
    'E': [0x7F,0x49,0x49,0x49,0x41],
    'F': [0x7F,0x09,0x09,0x09,0x01],
    'G': [0x3E,0x41,0x49,0x49,0x7A],
    'H': [0x7F,0x08,0x08,0x08,0x7F],
    'I': [0x41,0x7F,0x41,0x00,0x00],
    'L': [0x7F,0x40,0x40,0x40,0x40],
    'M': [0x7F,0x02,0x0C,0x02,0x7F],
    'N': [0x7F,0x04,0x08,0x10,0x7F],
    'O': [0x3E,0x41,0x41,0x41,0x3E],
    'P': [0x7F,0x09,0x09,0x09,0x06],
    'R': [0x7F,0x09,0x19,0x29,0x46],
    'S': [0x46,0x49,0x49,0x49,0x31],
    'T': [0x01,0x01,0x7F,0x01,0x01],
    'U': [0x3F,0x40,0x40,0x40,0x3F],
    'V': [0x1F,0x20,0x40,0x20,0x1F],
    'W': [0x3F,0x40,0x38,0x40,0x3F],
    'X': [0x63,0x14,0x08,0x14,0x63],
    'Y': [0x07,0x08,0x70,0x08,0x07],
    'Z': [0x61,0x51,0x49,0x45,0x43],
    '0': [0x3E,0x51,0x49,0x45,0x3E],
    '1': [0x00,0x42,0x7F,0x40,0x00],
    '2': [0x42,0x61,0x51,0x49,0x46],
    '3': [0x21,0x41,0x45,0x4B,0x31],
    '4': [0x18,0x14,0x12,0x7F,0x10],
    '5': [0x27,0x45,0x45,0x45,0x39],
    '6': [0x3C,0x4A,0x49,0x49,0x30],
    '7': [0x01,0x71,0x09,0x05,0x03],
    '8': [0x36,0x49,0x49,0x49,0x36],
    '9': [0x06,0x49,0x49,0x29,0x1E],
    '.': [0x00,0x60,0x60,0x00,0x00],
    ':': [0x00,0x36,0x36,0x00,0x00],
    '-': [0x08,0x08,0x08,0x08,0x08],
    '%': [0x23,0x13,0x08,0x64,0x62],
    'V': [0x1F,0x20,0x40,0x20,0x1F],
}

def oled_text(text, page, col=0):
    oled_cmd(0xB0 + page)
    oled_cmd(col & 0x0F)
    oled_cmd(0x10 | (col >> 4))
    for ch in text.upper():
        glyph = FONT.get(ch, [0x00]*5)
        i2c.writeto(OLED_ADDR, bytes([0x40] + glyph + [0x00]))

# ── MPU6050 ────────────────────────────────
MPU_ADDR = 0x68

def mpu6050_init():
    try:
        i2c.writeto_mem(MPU_ADDR, 0x6B, b'\x00')
        print("MPU6050 OK")
    except:
        print("MPU6050 not found")

def mpu6050_read():
    try:
        data = i2c.readfrom_mem(MPU_ADDR, 0x3B, 14)
        ax = round((data[0] << 8 | data[1]) / 16384.0, 2)
        ay = round((data[2] << 8 | data[3]) / 16384.0, 2)
        az = round((data[4] << 8 | data[5]) / 16384.0, 2)
        return ax, ay, az
    except:
        return 0.0, 0.0, 1.0

# ── Battery ────────────────────────────────
def read_battery():
    raw = adc.read_u16()
    voltage = round((raw / 65535) * 4.2, 2)
    percent = round(((voltage - 3.0) / 1.2) * 100, 1)
    percent = max(0, min(100, percent))
    return voltage, percent

# ── Atmosphere Simulation ──────────────────
def simulate_atmosphere(alt):
    if alt < 11000:
        temp = round(15.0 - (alt * 0.0065), 1)
    else:
        temp = -56.5
    pressure = round(101325 * math.exp(-alt / 8500), 1)
    humidity = round(max(0, 60 - (alt / 200)), 1)
    return temp, pressure, humidity

# ── GPS Mock ───────────────────────────────
def mock_gps(alt):
    lat = round(30.0444 + (alt * 0.00001), 4)
    lon = round(31.2357 + (alt * 0.000005), 4)
    return lat, lon

# ── Main ───────────────────────────────────
time.sleep(2)
mpu6050_init()

try:
    oled_init()
    oled_clear()
    oled_text("HAB LOGGER", 0)
    oled_text("STARTING...", 2)
    print("OLED OK")
except:
    print("OLED not found")

time.sleep(1)
altitude_sim = 0.0
time_sec = 0

print("=" * 50)
print("  HAB - High Altitude Balloon Data Logger")
print("=" * 50)

while True:
    altitude_sim += 50
    time_sec += 2

    temp, pressure, humidity = simulate_atmosphere(altitude_sim)
    ax, ay, az = mpu6050_read()
    voltage, batt_pct = read_battery()
    lat, lon = mock_gps(altitude_sim)

    # Update OLED
    try:
        oled_clear()
        oled_text(f"ALT: {altitude_sim}M", 0)
        oled_text(f"TMP: {temp}C", 1)
        oled_text(f"PRS: {pressure}PA", 2)
        oled_text(f"BAT: {voltage}V {batt_pct}%", 3)
        oled_text(f"AX:{ax} AZ:{az}", 4)
        oled_text(f"GPS:{lat}", 5)
    except:
        pass

    # Serial log
    print(f"T:{time_sec}s | Alt:{altitude_sim}m | Tmp:{temp}C | Prs:{pressure}Pa | Hum:{humidity}% | Ax:{ax} Ay:{ay} Az:{az} | Bat:{voltage}V({batt_pct}%) | GPS:{lat},{lon}")

    time.sleep(2)