import os

ROOT_DIR = r"C:\Users\LENOVO\Desktop\senti_ai"

# 1. Update train_superpacks_ai.py to use correct input dimensions and remove slice constraints
train_script = os.path.join(ROOT_DIR, "scratch", "train_superpacks_ai.py")
with open(train_script, "r", encoding="utf-8") as f:
    code = f.read()

# Fix classes input dims
code = code.replace("class TaxComplianceClassifier(nn.Module):\n    def __init__(self, input_dim=10, hidden_dim=32):", "class TaxComplianceClassifier(nn.Module):\n    def __init__(self, input_dim=11, hidden_dim=32):")
code = code.replace("class InsuranceNeedsAssessor(nn.Module):\n    def __init__(self, input_dim=8, hidden_dim=24):", "class InsuranceNeedsAssessor(nn.Module):\n    def __init__(self, input_dim=9, hidden_dim=24):")
code = code.replace("class AccountingBookkeepingClassifier(nn.Module):\n    def __init__(self, input_dim=8, hidden_dim=24):", "class AccountingBookkeepingClassifier(nn.Module):\n    def __init__(self, input_dim=9, hidden_dim=24):")
code = code.replace("class BankingTransactionScorer(nn.Module):\n    def __init__(self, input_dim=8, hidden_dim=32):", "class BankingTransactionScorer(nn.Module):\n    def __init__(self, input_dim=9, hidden_dim=32):")
code = code.replace("class WealthPortfolioAllocator(nn.Module):\n    def __init__(self, input_dim=8, hidden_dim=24):", "class WealthPortfolioAllocator(nn.Module):\n    def __init__(self, input_dim=9, hidden_dim=24):")

# Fix training instantiations input dims
code = code.replace("model_tax = TaxComplianceClassifier(input_dim=10)", "model_tax = TaxComplianceClassifier(input_dim=11)")
code = code.replace("model_ins = InsuranceNeedsAssessor(input_dim=8)", "model_ins = InsuranceNeedsAssessor(input_dim=9)")
code = code.replace("model_acc = AccountingBookkeepingClassifier(input_dim=8)", "model_acc = AccountingBookkeepingClassifier(input_dim=9)")
code = code.replace("model_bank = BankingTransactionScorer(input_dim=8)", "model_bank = BankingTransactionScorer(input_dim=9)")
code = code.replace("model_wealth = WealthPortfolioAllocator(input_dim=8)", "model_wealth = WealthPortfolioAllocator(input_dim=9)")

# Remove the slicing [:10] from tax features in training
code = code.replace("            ][:10]\n            X_tax.append(feat)", "\n            X_tax.append(feat)")

with open(train_script, "w", encoding="utf-8") as f:
    f.write(code)
print("Updated train_superpacks_ai.py input dimensions.")

# 2. Update model.py files in superpacks directories
packs = {
    "sentitax": {"dim": 11, "cls": "TaxComplianceClassifier"},
    "sentiinsurance": {"dim": 9, "cls": "InsuranceNeedsAssessor"},
    "sentiaccounting": {"dim": 9, "cls": "AccountingBookkeepingClassifier"},
    "sentibanking": {"dim": 9, "cls": "BankingTransactionScorer"},
    "sentiwealth": {"dim": 9, "cls": "WealthPortfolioAllocator"}
}

for name, info in packs.items():
    model_file = os.path.join(ROOT_DIR, name, "model.py")
    with open(model_file, "r", encoding="utf-8") as f:
        mcode = f.read()
    
    # Replace default input_dim in constructor
    mcode = mcode.replace(f"__init__(self, input_dim=10,", f"__init__(self, input_dim={info['dim']},")
    mcode = mcode.replace(f"__init__(self, input_dim=8,", f"__init__(self, input_dim={info['dim']},")
    
    with open(model_file, "w", encoding="utf-8") as f:
        f.write(mcode)
    print(f"Updated {name}/model.py input_dim to {info['dim']}.")