import cv2
import os
import numpy as np
import random

# -----------------------------------------
# PATHS
# -----------------------------------------
pos_dir = r"F:\My_new_folder\OccludedPedestrianDataset\positive_samples"
neg_dir = r"F:\My_new_folder\OccludedPedestrianDataset\negative_samples"

weights_float_path = r"F:\My_new_folder\OccludedPedestrianDataset\svm_float_weights_2c.txt"
weights_fixed_path = r"F:\My_new_folder\OccludedPedestrianDataset\svm_fixed_weights_2c.txt"

# -----------------------------------------
# CUSTOM HOG (1152 FEATURE VERSION)
# -----------------------------------------
hog = cv2.HOGDescriptor(
    _winSize=(64, 128),
    _blockSize=(16, 16),
    _blockStride=(16, 16),
    _cellSize=(8, 8),
    _nbins=9
)

def compute_hog(image_path):
    img = cv2.imread(image_path)
    img = cv2.resize(img, (64,128))
    feat = hog.compute(img).flatten()

    if len(feat) != 1152:
        print("ERROR: HOG feature size =", len(feat))
        exit()

    return feat

# -----------------------------------------
# LOAD DATASET
# -----------------------------------------
X = []
y = []

for f in os.listdir(pos_dir):
    if f.endswith(".png"):
        X.append(compute_hog(os.path.join(pos_dir, f)))
        y.append(1)

for f in os.listdir(neg_dir):
    if f.endswith(".png"):
        X.append(compute_hog(os.path.join(neg_dir, f)))
        y.append(0)

X = np.array(X, dtype=np.float32)
y = np.array(y, dtype=np.int32)

print("Dataset size:", len(X))
print("Feature length =", X[0].shape[0])   # MUST PRINT 1152

# -----------------------------------------
# 80/20 SPLIT
# -----------------------------------------
random.seed(42)
indices = list(range(len(X)))
random.shuffle(indices)

split = int(0.8 * len(X))
train_idx = indices[:split]
test_idx  = indices[split:]

X_train, y_train = X[train_idx], y[train_idx]
X_test,  y_test  = X[test_idx],  y[test_idx]

print("TRAIN =", len(X_train), "TEST =", len(X_test))

# -----------------------------------------
# TRAIN SVM
# -----------------------------------------
svm = cv2.ml.SVM_create()
svm.setType(cv2.ml.SVM_C_SVC)
svm.setKernel(cv2.ml.SVM_LINEAR)
svm.setC(0.05)

svm.train(X_train, cv2.ml.ROW_SAMPLE, y_train)

rho, alpha, svidx = svm.getDecisionFunction(0)
w = -alpha @ svm.getSupportVectors()
w = w.flatten()

print("WEIGHT VECTOR SIZE =", len(w))  # MUST BE 1152
print("BIAS =", rho)

# -----------------------------------------
# SAVE FLOAT WEIGHTS
# -----------------------------------------
with open(weights_float_path, "w") as f:
    weights_str = ", ".join([f"{v}" for v in w])
    f.write(weights_str + ", " + str(rho))

# -----------------------------------------
# SAVE FIXED WEIGHTS
# -----------------------------------------
SCALE = 256  # Q8.8

with open(weights_fixed_path, "w") as f:
    weights_fixed = ", ".join([str(int(v * SCALE)) for v in w])
    f.write(weights_fixed + ", " + str(int(rho * SCALE)))

# -----------------------------------------
# TEST
# -----------------------------------------
def svm_predict(x, w, b):
    return 1 if np.dot(w, x) + b > 0 else 0

TP = FP = TN = FN = 0

for feat, label in zip(X_test, y_test):
    pred = svm_predict(feat, w, rho)

    if pred == 1 and label == 1: TP += 1
    if pred == 1 and label == 0: FP += 1
    if pred == 0 and label == 0: TN += 1
    if pred == 0 and label == 1: FN += 1

accuracy  = (TP + TN) / (TP + TN + FP + FN + 1e-6)
precision = TP / (TP + FP + 1e-6)
recall    = TP / (TP + FN + 1e-6)
f1_score  = 2 * precision * recall / (precision + recall + 1e-6)

print("\n=== CONFUSION MATRIX ===")
print("TP =", TP)
print("FP =", FP)
print("TN =", TN)
print("FN =", FN)

print("\nACCURACY  =", accuracy)
print("PRECISION =", precision)
print("RECALL    =", recall)
print("F1 SCORE  =", f1_score)