import numpy as np
import torch
from scipy.signal import butter, filtfilt
import secrets

def numpy_to_storage(labels, data, storage_file, datatype=None):
    assert data.shape[1] == len(labels), "# labels doesn't match columns"
    assert labels[0] == "time"

    f = open(storage_file, 'w')
    # Old style
    if datatype is None:
        f = open(storage_file, 'w')
        f.write('name %s\n' % storage_file)
        f.write('datacolumns %d\n' % data.shape[1])
        f.write('datarows %d\n' % data.shape[0])
        f.write('range %f %f\n' % (np.min(data[:, 0]), np.max(data[:, 0])))
        f.write('endheader \n')
    # New style
    else:
        if datatype == 'IK':
            f.write('Coordinates\n')
        elif datatype == 'ID':
            f.write('Inverse Dynamics Generalized Forces\n')
        elif datatype == 'GRF':
            f.write('%s\n' % storage_file)
        elif datatype == 'muscle_forces':
            f.write('ModelForces\n')
        f.write('version=1\n')
        f.write('nRows=%d\n' % data.shape[0])
        f.write('nColumns=%d\n' % data.shape[1])
        if datatype == 'IK':
            f.write('inDegrees=yes\n\n')
            f.write('Units are S.I. units (second, meters, Newtons, ...)\n')
            f.write(
                "If the header above contains a line with 'inDegrees', this indicates whether rotational values are in degrees (yes) or radians (no).\n\n")
        elif datatype == 'ID':
            f.write('inDegrees=no\n')
        elif datatype == 'GRF':
            f.write('inDegrees=yes\n')
        elif datatype == 'muscle_forces':
            f.write('inDegrees=yes\n\n')
            f.write('This file contains the forces exerted on a model during a simulation.\n\n')
            f.write("A force is a generalized force, meaning that it can be either a force (N) or a torque (Nm).\n\n")
            f.write('Units are S.I. units (second, meters, Newtons, ...)\n')
            f.write('Angles are in degrees.\n\n')

        f.write('endheader \n')

    for i in range(len(labels)):
        f.write('%s\t' % labels[i])
    f.write('\n')

    for i in range(data.shape[0]):
        for j in range(data.shape[1]):
            f.write('%20.8f\t' % data[i, j])
        f.write('\n')

    f.close()


def lowpass_filter(data, cutoff_cycles=2, num_samples=24, order=4):
    nyquist = num_samples / 2  # Max frequency is Nyquist (12 cycles/stride for 24 samples)
    normal_cutoff = cutoff_cycles / nyquist  # Convert cycles per stride to normalized frequency
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    return filtfilt(b, a, data, axis=0)

# Define a function to get a random stride of the dataloader, save it to a mot file and print the metadata as well as the reconstruction
def save_random_stride(model, dataloader, joints, savepath, savepath_recon, device,
                       from_all_batches=False, filter=False):
    # ----- choose a random example using secrets -----
    if from_all_batches:
        batch_data, batch_metadata = [], []
        for batch in dataloader:
            batch_data.append(batch["features"])
            batch_metadata.extend(batch["metadata"])
        batch_data = torch.cat(batch_data, dim=0)

        stride_idx = secrets.randbelow(batch_data.shape[0])  # cryptographically strong
        stride_data = batch_data[stride_idx]
        stride_metadata = batch_metadata[stride_idx]
    else:
        ds = dataloader.dataset
        ds_idx = secrets.randbelow(len(ds))
        item = ds[ds_idx]  # expects {"features": (S,T,D), "metadata": ...}

        feats = item["features"]                 # (num_strides, T, D)
        s_idx = secrets.randbelow(feats.shape[0])
        stride_data = feats[s_idx]

        meta = item["metadata"]
        try:
            stride_metadata = meta.iloc[s_idx]   # per-stride metadata (DataFrame)
        except Exception:
            stride_metadata = meta               # per-subject metadata

    print(f"Stride Data Shape: {stride_data.shape}")

    # ----- forward pass -----
    model.eval()
    with torch.no_grad():
        reconstructed_data, _, mu, _, _ = model(stride_data.unsqueeze(0).to(device))
        reconstructed_data = reconstructed_data.view(-1, 24, 32).cpu().numpy().squeeze()
        stride_data = stride_data.cpu().numpy().squeeze()

    # ----- keep requested joints -----
    stride_data = stride_data[:, :len(joints)]
    reconstructed_data = reconstructed_data[:, :len(joints)]

    # ----- sampling freq -----
    stride_time = stride_data[:, 0]
    fs = 1.0 / np.mean(np.diff(stride_time))
    print(f"Sampling Frequency: {fs} Hz")

    # ----- optional filtering -----
    if filter:
        reconstructed_data = lowpass_filter(stride_data, cutoff_cycles=2, num_samples=24)

    reconstructed_data[:, 0] = stride_time
    joints = [j.replace('_ips', '_r').replace('_contra', '_l') for j in joints]

    numpy_to_storage(joints, stride_data, savepath, datatype='IK')
    numpy_to_storage(joints, reconstructed_data, savepath_recon, datatype='IK')

    print(f"Stride Metadata: {stride_metadata}")
    return mu  # torch tensor on device