import os
import numpy as np
import matplotlib.pyplot as plt

# ============ User settings ============
caseID = 1        # 1=FC, 2=MC, 3=NC
wallID = 1        # 1=hotWall, 2=baffleLeft, 3=baffleRight, 4=coldWall
varID  = 3        # 1=T, 2=q_w, 3=tau_w
# ========================================

caseNames = ['FC', 'MC', 'NC']
caseName = caseNames[caseID - 1]

wallNames = ['hotWall', 'baffleLeft', 'baffleRight', 'coldWall']
wallName = wallNames[wallID - 1]

varNames = ['T', 'q_w', 'tau_w']
varUnits = ['K', 'W/m^2', 'Pa']
varCol = varID  # column 0 is y, so column index = varID (1=T,...)

varName = varNames[varID - 1]
varUnit = varUnits[varID - 1]

baseDir = os.path.dirname(os.path.abspath(__file__))
caseDir = os.path.join(baseDir, caseName)

fig, ax = plt.subplots()

# ---- DNS ----
dnsFile = os.path.join(caseDir, 'DNS', f'{wallName}.dat')
if os.path.isfile(dnsFile):
    data = np.loadtxt(dnsFile, comments='#')
    ax.plot(data[:, 0], data[:, varCol], 'k-', linewidth=2, label='DNS')

# ---- RANS models (every subfolder of RANS/ is one model) ----
ransDir = os.path.join(caseDir, 'RANS')
models = sorted(m for m in os.listdir(ransDir) if os.path.isdir(os.path.join(ransDir, m)))

colors = plt.cm.tab10(np.linspace(0, 1, 10))
markers = ['o', 's', '^', 'v', 'd', 'p', 'h', 'x', '+', '*']

for i, model in enumerate(models):
    f = os.path.join(ransDir, model, f'{wallName}.dat')
    if not os.path.isfile(f):
        continue
    data = np.loadtxt(f, comments='#')
    mk = markers[i % len(markers)]
    nm = max(1, round(len(data) / 20))

    ax.plot(data[:, 0], data[:, varCol], '-', color=colors[i % len(colors)],
            marker=mk, markevery=nm, markersize=5,
            label=model.replace('_', ' '))

ax.set_xlabel('y (m)')
ax.set_ylabel(f'{varName} ({varUnit})')
ax.set_title(f'{caseName} - {wallName}')
ax.grid(True)
ax.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))

fig.tight_layout()
plt.show()