import yaml
import io


def mwdw_to_dat(fn_mix, n_weeks, isstringio=True):
    """
    Convert multi-week days worked inputs into GMPL dat form.

    :param fn_mix: filename of yaml tour type mix file
    :param n_weeks:
    :param isstringio: true to return StringIO object, false to return string
    :return: GMPL dat code for mwdw patterns either as a StringIO
        object or a string.

    Example:
        param num_mwdw_patterns:=
        1 1
        2 2
        3 2
        4 2
        5 1
        6 2
        7 2
        8 4
        ;

        param A_mwdw:=
        1 1 1 5
        1 1 2 5
        1 1 3 5
        """

    f_mix_in = open(fn_mix, "r")
    ttspec = yaml.safe_load(f_mix_in)
    f_mix_in.close()

    # n_ttypes = np.size(ttspec['tourtypes'])
    mwdw_rows = []
    num_mwdw_rows = []
    for m in ttspec['tourtypes']:
        num_rows = 0
        for p in range(len(m['multiweek_num_days_worked'])):
            num_rows += 1
            for w in range(n_weeks):
                mwdw_row = [m['ttnum'], p + 1, w + 1, m['multiweek_num_days_worked'][p][w]]
                mwdw_rows.append(mwdw_row)

        temp_list = [m['ttnum'], num_rows]
        num_mwdw_rows.append(temp_list)

    param = 'param num_mwdw_patterns:=\n'
    for val in num_mwdw_rows:
        data_row = ' '.join(map(str, val)) + '\n'
        param += data_row
    param += ";\n"

    param += '\nparam A_mwdw:=\n'
    for val in mwdw_rows:
        data_row = ' '.join(map(str, val)) + '\n'
        param += data_row
    param += ";\n"

    if isstringio:
        param_out = io.StringIO()
        param_out.write(param)
        return param_out.getvalue()
    else:
        return param