"""Utilities employed throughout the package."""

import numpy as np


__all__ = ["complex2real", "save_xy", "print_color"]


def complex2real(x: np.ndarray) -> np.ndarray:
    """
    Split the real and imaginary parts of complex array.

    NOTE:
    1. The format of output is
        x[:,0].real, x[:,0].imag, x[:,1].real, x[:,1].imag, ...
    2. If x is a real matrix, then a copy of x is returned.

    :param x: (num_row, num_cols) complex128 array
    :return: (num_rows, 2*num_cols) float64 array if x is complex array,
        otherwise (num_rows, num_cols)
    """
    # Convert x into matrix if it is a ROW vector
    if len(x.shape) == 1:
        x = x[:, np.newaxis]

    # Copy real and imaginary parts of x
    num_rows, num_cols = x.shape
    if isinstance(x.item(0, 0), complex):
        extended_data = np.zeros((num_rows, 2*num_cols), dtype=np.float64)
        extended_data[:, 0::2] = x.real
        extended_data[:, 1::2] = x.imag
    else:
        extended_data = x.copy()
    return extended_data


def save_xy(file_name: str, x: np.ndarray, y: np.ndarray) -> None:
    """
    Save vectors x and y in two columns in one file.

    :param file_name: name of output file
    :param x: (num_data, ) float64 array
    :param y: (num_data, ) float64 array
    """
    x_ext = complex2real(x)
    y_ext = complex2real(y)
    data_all = np.hstack((x_ext, y_ext))
    np.savetxt(file_name, data_all)


def print_color(text: str, color: str = "w") -> None:
    """
    Print colored text.

    :param text: text to print
    :param color: color of text
    :returns: None
    """
    code = {"r": 31, "g": 32, "y": 33, "b": 34, "m": 35, "c": 36, "w": 37}
    msg = f"\033[{code[color]}m{text}\033[0m"
    print(msg)
