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

# ============ User settings ============
caseID = 1        # 1=FC, 2=MC, 3=NC
lineID = 4        # horizontal profile line number, 1-12
varID  = 2        # 1=U, 2=V, 3=k, 4=T
# ========================================

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

varNames = ['U', 'V', 'k', 'T']
varUnits = ['m/s', 'm/s', 'm^2/s^2', 'K']
varCol = varID  # column 0 is x, so column index = varID (1=U,...)

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'line{lineID}.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'line{lineID}.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('x (m)')
ax.set_ylabel(f'{varName} ({varUnit})')
ax.set_title(f'{caseName} - line {lineID}')
ax.grid(True)
ax.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))

fig.tight_layout()
plt.show()