import torch


def secondary_sequence_metrics(x, ss_num=7):
    """
    Compute metrics associated with with secondary structure.

    It counts how many times a secondary structure appears in the proteins, and the median length of a sequence of
    secondary structure, e.g. H, H, ..., H.

    Parameters
    ----------
    x : torch.Tensor:
        The node features.
    ss_num : int, optional
        The number of secondary structures. The default is 7.

    Returns
    -------
    scores : torch.Tensor
        The score having the following format: [SS_1_count, ..., SS_ss_num_count, median_seq_len]
    """
    # Get number of amino acids
    n = x.shape[0]
    # Initialize scores
    ss_counts, seq_len = [0]*ss_num, []
    current_seq_len, current_ss = 0, -1
    # Compute the scores
    for i in range(n):
        ss = int(x[i,1])
        ss_counts[ss-1] += 1
        if current_ss == -1:
            current_seq_len += 1
            current_ss = ss
        elif ss != current_ss:
            seq_len.append(current_seq_len)
            current_seq_len = 1
            current_ss = ss
        else:
            current_seq_len += 1
    seq_len.append(current_seq_len) # last one to add
    ss_counts = [x/n for x in ss_counts]
    seq_len = float(torch.median(torch.tensor(seq_len)))
    scores = torch.tensor(ss_counts + [seq_len])
    return scores